[3.15] gh-160522: Restore CJK encoder state after a failed incremental encode (GH-160523) (#161000)

gh-160522: Restore CJK encoder state after a failed incremental encode (GH-160523)

A failed IncrementalEncoder.encode() or StreamWriter.write() of a
stateful CJK codec kept the shift state reached before the error
although its output was discarded.
(cherry picked from commit 5b25582)

Co-authored-by: Gregory P. Smith <68491+gpshead@users.noreply.github.com>

Author: miss-islington

Lib/test/test_multibytecodec.py

@@ -438,6 +438,48 @@ def test_incrementalencoder_final(self):
         self.assertEqual(output, self.expected_reset)
         self.assertEqual(encoder.encode('', final=True), b'')
 
+    def test_incrementalencoder_error(self):
+        # The output of a failed encode() is lost, so the shift state
+        # must not be changed.
+        encoder = codecs.getincrementalencoder(self.encoding)()
+        state = encoder.getstate()
+        for final in False, True:
+            self.assertRaises(UnicodeEncodeError, encoder.encode,
+                              self.text + '\U0001f600', final)
+            self.assertEqual(encoder.getstate(), state)
+        self.assertEqual(encoder.encode(self.text), self.expected)
+        state = encoder.getstate()
+        for final in False, True:
+            self.assertRaises(UnicodeEncodeError, encoder.encode,
+                              'a\U0001f600', final)
+            self.assertEqual(encoder.getstate(), state)
+        self.assertEqual(encoder.encode('a'), self.reset + b'a')
+
+    def test_streamwriter_error(self):
+        stream = io.BytesIO()
+        writer = codecs.getwriter(self.encoding)(stream)
+        self.assertRaises(UnicodeEncodeError, writer.write,
+                          self.text + '\U0001f600')
+        self.assertEqual(stream.getvalue(), b'')
+        writer.write(self.text)
+        self.assertEqual(stream.getvalue(), self.expected)
+
+    def test_textiowrapper_error(self):
+        raw = io.BytesIO()
+        with io.TextIOWrapper(raw, encoding=self.encoding) as f:
+            self.assertRaises(UnicodeEncodeError, f.write,
+                              self.text + '\U0001f600')
+            f.write(self.text)
+            f.flush()
+            self.assertEqual(raw.getvalue(), self.expected)
+
+class TestISO2022KRStateful(TestStateful):
+    text = '\uac00\uac00'
+    encoding = 'iso-2022-kr'
+    expected = b'\x1b$)C\x0e0!0!'
+    reset = b'\x0f'
+    expected_reset = expected + reset
+
 class TestHZStateful(TestStateful):
     text = '\u804a\u804a'
     encoding = 'hz'

Misc/NEWS.d/next/Library/2026-10-22-08-18-31.gh-issue-160522.cJkSt8.rst

@@ -0,0 +1,4 @@
+Fix the incremental encoders and stream writers of the stateful CJK codecs
+(``iso2022_*`` and ``hz``) keeping a changed shift state after an
+:meth:`!encode` or :meth:`!write` call failed.  The next call could omit the
+escape sequence, writing multibyte characters in ASCII mode.

Modules/cjkcodecs/multibytecodec.c

@@ -793,6 +793,7 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
     PyObject *inbuf = NULL;
     Py_ssize_t inpos, datalen;
     PyObject *origpending = NULL;
+    MultibyteCodec_State origstate;
 
     if (PyUnicode_Check(unistr))
         ucvt = NULL;
@@ -808,6 +809,7 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
         }
     }
 
+    origstate = ctx->state;
     if (ctx->pending) {
         PyObject *inbuf_tmp;
 
@@ -865,6 +867,9 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
     return r;
 
 errorexit:
+    /* recover the original codec state: the discarded output may have
+     * included an escape sequence for a state change */
+    ctx->state = origstate;
     Py_XDECREF(r);
     Py_XDECREF(ucvt);
     Py_XDECREF(origpending);