gh-154859: Keep the iconv shift state across incremental decode calls

Include/internal/pycore_unicodeobject.h

@@ -185,12 +185,19 @@ extern int _PyUnicodeWriter_FormatV(
 /* --- iconv Codec -------------------------------------------------------- */
 
 #ifdef HAVE_ICONV
+#include <iconv.h>
+
+/* Open a conversion for decoding ENCODING, to reuse across calls.  Returns
+   (iconv_t)-1 with an exception set on failure. */
+extern iconv_t _PyUnicode_IconvOpenDecoder(const char *encoding);
+
 extern PyObject* _PyUnicode_DecodeIconv(
     const char *encoding,       /* iconv encoding name */
     const char *string,         /* encoded string */
     Py_ssize_t length,          /* size of string */
     const char *errors,         /* error handling */
-    Py_ssize_t *consumed);      /* bytes consumed, or NULL for non-stateful */
+    Py_ssize_t *consumed,       /* bytes consumed, or NULL for non-stateful */
+    iconv_t *cdp);              /* conversion to reuse, or NULL to open one */
 
 extern PyObject* _PyUnicode_EncodeIconv(
     const char *encoding,       /* iconv encoding name */

Lib/encodings/_iconv_codecs.py

@@ -1,7 +1,7 @@
 import codecs
 
 def create_iconv_codec(name, encoding):
-    from _codecs import iconv_encode, iconv_decode
+    from _codecs import iconv_encode, iconv_decode, iconv_state
 
     def encode(input, errors='strict'):
         return iconv_encode(encoding, input, errors)
@@ -14,16 +14,32 @@ def encode(self, input, final=False):
             return iconv_encode(encoding, input, self.errors)[0]
 
     class IncrementalDecoder(codecs.BufferedIncrementalDecoder):
+        def __init__(self, errors='strict'):
+            super().__init__(errors)
+            self._state = iconv_state(encoding)
+
         def _buffer_decode(self, input, errors, final):
-            return iconv_decode(encoding, input, errors, final)
+            return iconv_decode(encoding, input, errors, final, self._state)
+
+        def reset(self):
+            super().reset()
+            self._state = iconv_state(encoding)
 
     class StreamWriter(codecs.StreamWriter):
         def encode(self, input, errors='strict'):
             return iconv_encode(encoding, input, errors)
 
     class StreamReader(codecs.StreamReader):
+        def __init__(self, stream, errors='strict'):
+            super().__init__(stream, errors)
+            self._state = iconv_state(encoding)
+
         def decode(self, input, errors, final=False):
-            return iconv_decode(encoding, input, errors, final)
+            return iconv_decode(encoding, input, errors, final, self._state)
+
+        def reset(self):
+            super().reset()
+            self._state = iconv_state(encoding)
 
     return codecs.CodecInfo(
         name=name,

Lib/test/test_codecs.py

@@ -3676,6 +3676,10 @@ def iconv_encoding_available(name):
     ('ISO-8859-1', 'Grüße'),
 ]
 _ICONV_MULTIBYTE = ['EUC-JP', 'SHIFT_JIS', 'GBK', 'GB18030', 'BIG5']
+# Stateful encodings: the shift state set by an escape sequence has to survive
+# from one incremental call to the next.  CPython has no built-in codec for
+# these, so the plain name reaches the iconv codec.
+_ICONV_STATEFUL = ['ISO-2022-CN']
 # Encodings iconv may provide but for which CPython has no built-in codec
 # (cp1047 is EBCDIC, i.e. not ASCII-compatible).
 _ICONV_ONLY = ['cp1047', 'cp1133', 'GEORGIAN-PS', 'ARMSCII-8']
@@ -3804,6 +3808,26 @@ def test_stream(self):
         reader = codecs.getreader('iconv:' + enc)(io.BytesIO(raw))
         self.assertEqual(reader.read(), text)
 
+    def test_incremental_decode_shift_state(self):
+        enc = self.require(*_ICONV_STATEFUL)
+        text = 'ABC\u4e2d\u6587DEF'
+        data = codecs.encode(text, 'iconv:' + enc)
+        self.assertEqual(codecs.decode(data, 'iconv:' + enc), text)
+        dec = codecs.getincrementaldecoder('iconv:' + enc)()
+        out = ''.join(dec.decode(data[i:i+1]) for i in range(len(data)))
+        out += dec.decode(b'', True)
+        self.assertEqual(out, text)
+        # reset() starts a new conversion, so decoding can begin again.
+        dec.reset()
+        self.assertEqual(dec.decode(data, True), text)
+
+    def test_stream_shift_state(self):
+        enc = self.require(*_ICONV_STATEFUL)
+        text = 'ABC\u4e2d\u6587DEF'
+        raw = codecs.encode(text, 'iconv:' + enc)
+        reader = codecs.getreader('iconv:' + enc)(io.BytesIO(raw))
+        self.assertEqual(''.join(iter(lambda: reader.read(1), '')), text)
+
     def test_encode_kinds(self):
         # The string's own buffer is fed to iconv per storage kind; check each
         # of the 1-, 2- and 4-byte kinds against the built-in codec.

Modules/_codecsmodule.c

@@ -644,24 +644,89 @@ _codecs_code_page_decode_impl(PyObject *module, int codepage,
 
 #ifdef HAVE_ICONV
 
+#ifdef HAVE_ICONV
+#define ICONV_STATE_CAPSULE "_codecs.iconv_state"
+
+static void
+iconv_state_destructor(PyObject *capsule)
+{
+    iconv_t *cdp = PyCapsule_GetPointer(capsule, ICONV_STATE_CAPSULE);
+    if (cdp == NULL) {
+        PyErr_Clear();
+        return;
+    }
+    iconv_close(*cdp);
+    PyMem_Free(cdp);
+}
+#endif
+
+/*[clinic input]
+_codecs.iconv_state
+
+    encoding: str
+    /
+
+Open an iconv conversion for decoding, to reuse across calls.
+
+The result is an opaque object.  Reusing one conversion keeps the
+shift state of a stateful encoding, such as ISO-2022-CN, across calls.
+[clinic start generated code]*/
+
+static PyObject *
+_codecs_iconv_state_impl(PyObject *module, const char *encoding)
+/*[clinic end generated code: output=4100a9b65a64d65c input=6ffbfa5bb1d2d208]*/
+{
+#ifdef HAVE_ICONV
+    iconv_t *cdp = PyMem_Malloc(sizeof(iconv_t));
+    if (cdp == NULL) {
+        return PyErr_NoMemory();
+    }
+    *cdp = _PyUnicode_IconvOpenDecoder(encoding);
+    if (*cdp == (iconv_t)-1) {
+        PyMem_Free(cdp);
+        return NULL;
+    }
+    PyObject *capsule = PyCapsule_New(cdp, ICONV_STATE_CAPSULE,
+                                      iconv_state_destructor);
+    if (capsule == NULL) {
+        iconv_close(*cdp);
+        PyMem_Free(cdp);
+        return NULL;
+    }
+    return capsule;
+#else
+    PyErr_SetString(PyExc_LookupError, "iconv is not available");
+    return NULL;
+#endif
+}
+
 /*[clinic input]
 _codecs.iconv_decode
     encoding: str
     data: Py_buffer
     errors: str(accept={str, NoneType}) = None
     final: bool = False
+    state: object = None
     /
 [clinic start generated code]*/
 
 static PyObject *
 _codecs_iconv_decode_impl(PyObject *module, const char *encoding,
-                          Py_buffer *data, const char *errors, int final)
-/*[clinic end generated code: output=6c6145a9decc2ba8 input=d15a04d7d3a3e0cd]*/
+                          Py_buffer *data, const char *errors, int final,
+                          PyObject *state)
+/*[clinic end generated code: output=ed99087a9b21d007 input=b23f4298c963f9c5]*/
 {
+    iconv_t *cdp = NULL;
+    if (state != Py_None) {
+        cdp = PyCapsule_GetPointer(state, ICONV_STATE_CAPSULE);
+        if (cdp == NULL) {
+            return NULL;
+        }
+    }
     Py_ssize_t consumed = data->len;
     PyObject *decoded = _PyUnicode_DecodeIconv(encoding, data->buf, data->len,
                                                errors,
-                                               final ? NULL : &consumed);
+                                               final ? NULL : &consumed, cdp);
     return codec_tuple(decoded, consumed);
 }
 
@@ -1155,6 +1220,7 @@ static PyMethodDef _codecs_functions[] = {
     _CODECS_CODE_PAGE_DECODE_METHODDEF
     _CODECS_ICONV_ENCODE_METHODDEF
     _CODECS_ICONV_DECODE_METHODDEF
+    _CODECS_ICONV_STATE_METHODDEF
     _CODECS_REGISTER_ERROR_METHODDEF
     _CODECS__UNREGISTER_ERROR_METHODDEF
     _CODECS_LOOKUP_ERROR_METHODDEF

Objects/unicodeobject.c

@@ -8239,26 +8239,42 @@ iconv_open_or_set_error(const char *tocode, const char *fromcode,
     return cd;
 }
 
+iconv_t
+_PyUnicode_IconvOpenDecoder(const char *encoding)
+{
+    return iconv_open_or_set_error(ICONV_PIVOT, encoding, encoding);
+}
+
 /*
  * Decode bytes with iconv() into a str.
  *
  * The input is converted to native-endian UTF-32 one chunk at a time and
  * appended to a _PyUnicodeWriter.  If *consumed* is non-NULL the decode is
  * stateful: a trailing incomplete sequence stops and sets *consumed*.
+ *
+ * If *cdp* is non-NULL the conversion it points to is used and left open, so
+ * that a shift state survives from one call to the next.
  */
 PyObject *
 _PyUnicode_DecodeIconv(const char *encoding,
                        const char *s, Py_ssize_t size,
-                       const char *errors, Py_ssize_t *consumed)
+                       const char *errors, Py_ssize_t *consumed,
+                       iconv_t *cdp)
 {
     if (size < 0) {
         PyErr_BadInternalCall();
         return NULL;
     }
 
-    iconv_t cd = iconv_open_or_set_error(ICONV_PIVOT, encoding, encoding);
-    if (cd == (iconv_t)-1) {
-        return NULL;
+    iconv_t cd;
+    if (cdp != NULL) {
+        cd = *cdp;
+    }
+    else {
+        cd = _PyUnicode_IconvOpenDecoder(encoding);
+        if (cd == (iconv_t)-1) {
+            return NULL;
+        }
     }
 
     /* Scratch buffer for one iconv() output chunk, as UTF-32 code points. */
@@ -8337,13 +8353,17 @@ _PyUnicode_DecodeIconv(const char *encoding,
     if (consumed != NULL) {
         *consumed = in - starts;
     }
-    iconv_close(cd);
+    if (cdp == NULL) {
+        iconv_close(cd);
+    }
     Py_XDECREF(errorHandler);
     Py_XDECREF(exc);
     return _PyUnicodeWriter_Finish(&writer);
 
 error:
-    iconv_close(cd);
+    if (cdp == NULL) {
+        iconv_close(cd);
+    }
     _PyUnicodeWriter_Dealloc(&writer);
     Py_XDECREF(errorHandler);
     Py_XDECREF(exc);