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, IconvDecoder
 
     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._decoder = IconvDecoder(encoding)
+
         def _buffer_decode(self, input, errors, final):
-            return iconv_decode(encoding, input, errors, final)
+            return self._decoder.decode(input, errors, final)
+
+        def reset(self):
+            super().reset()
+            self._decoder = IconvDecoder(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._decoder = IconvDecoder(encoding)
+
         def decode(self, input, errors, final=False):
-            return iconv_decode(encoding, input, errors, final)
+            return self._decoder.decode(input, errors, final)
+
+        def reset(self):
+            super().reset()
+            self._decoder = IconvDecoder(encoding)
 
     return codecs.CodecInfo(
         name=name,

Lib/test/test_codecs.py

@@ -3677,6 +3677,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']
@@ -3820,6 +3824,33 @@ def test_stream(self):
         reader = codecs.getreader('iconv:' + enc)(io.BytesIO(raw))
         self.assertEqual(reader.read(), text)
 
+    def require_stateful(self):
+        # Encoded here rather than by the platform: an iconv that provides the
+        # encoding may still be unable to encode the sample, as macOS and iOS
+        # cannot for ISO-2022-CN.  Decoding is what these tests are about, so
+        # the bytes are fixed and only decoding has to work.
+        enc = self.require(*_ICONV_STATEFUL)
+        text = 'ABC\u4e2d\u6587DEF'
+        data = b'ABC\x1b$)A\x0eVPND\x0fDEF\x0f'
+        if codecs.decode(data, 'iconv:' + enc) != text:
+            self.skipTest('%s: this iconv cannot decode the sample' % enc)
+        return enc, text, data
+
+    def test_incremental_decode_shift_state(self):
+        enc, text, data = self.require_stateful()
+        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, text, data = self.require_stateful()
+        reader = codecs.getreader('iconv:' + enc)(io.BytesIO(data))
+        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

@@ -32,16 +32,45 @@ Copyright (c) Corporation for National Research Initiatives.
 
 #include "Python.h"
 #include "pycore_codecs.h"        // _PyCodec_Lookup()
+#include "pycore_pymem.h"         // _PyMem_Strdup()
 #include "pycore_unicodeobject.h" // _PyUnicode_EncodeCharmap
 
 #ifdef MS_WINDOWS
 #include <windows.h>
 #endif
 
+typedef struct {
+    PyTypeObject *IconvDecoderType;
+} _codecs_state;
+
+static inline _codecs_state *
+get_codecs_state(PyObject *module)
+{
+    void *state = PyModule_GetState(module);
+    assert(state != NULL);
+    return (_codecs_state *)state;
+}
+
+static struct PyModuleDef codecsmodule;
+
+#define get_codecs_state_by_type(type) \
+    (get_codecs_state(PyType_GetModuleByDef(type, &codecsmodule)))
+
+#ifdef HAVE_ICONV
+typedef struct {
+    PyObject_HEAD
+    iconv_t cd;
+    char *encoding;
+} iconv_decoder_object;
+
+#define iconv_decoder_CAST(op)  ((iconv_decoder_object *)(op))
+#endif
+
 /*[clinic input]
 module _codecs
+class _codecs.IconvDecoder "iconv_decoder_object *" "get_codecs_state_by_type(type)->IconvDecoderType"
 [clinic start generated code]*/
-/*[clinic end generated code: output=da39a3ee5e6b4b0d input=e1390e3da3cb9deb]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=17a98f0bef095afe]*/
 
 #include "pycore_runtime.h"
 #include "clinic/_codecsmodule.c.h"
@@ -644,6 +673,108 @@ _codecs_code_page_decode_impl(PyObject *module, int codepage,
 
 #ifdef HAVE_ICONV
 
+/*[clinic input]
+@classmethod
+_codecs.IconvDecoder.__new__
+
+    encoding: str
+    /
+
+Decoder holding one iconv conversion, to reuse across calls.
+
+Reusing one conversion keeps the shift state of a stateful encoding,
+such as ISO-2022-CN, from one call to the next.
+[clinic start generated code]*/
+
+static PyObject *
+_codecs_IconvDecoder_impl(PyTypeObject *type, const char *encoding)
+/*[clinic end generated code: output=6e5181abedc4ae7c input=c53769050ff2b196]*/
+{
+    char *name = _PyMem_Strdup(encoding);
+    if (name == NULL) {
+        return PyErr_NoMemory();
+    }
+    iconv_t cd = _PyUnicode_IconvOpenDecoder(encoding);
+    if (cd == (iconv_t)-1) {
+        PyMem_Free(name);
+        return NULL;
+    }
+    iconv_decoder_object *self = (iconv_decoder_object *)type->tp_alloc(type, 0);
+    if (self == NULL) {
+        iconv_close(cd);
+        PyMem_Free(name);
+        return NULL;
+    }
+    self->cd = cd;
+    self->encoding = name;
+    return (PyObject *)self;
+}
+
+/*[clinic input]
+_codecs.IconvDecoder.decode
+
+    data: Py_buffer
+    errors: str(accept={str, NoneType}) = None
+    final: bool = False
+    /
+[clinic start generated code]*/
+
+static PyObject *
+_codecs_IconvDecoder_decode_impl(iconv_decoder_object *self, Py_buffer *data,
+                                 const char *errors, int final)
+/*[clinic end generated code: output=0e8812b6b422fc97 input=9bea42dd438d03af]*/
+{
+    Py_ssize_t consumed = data->len;
+    PyObject *decoded = _PyUnicode_DecodeIconv(self->encoding, data->buf,
+                                               data->len, errors,
+                                               final ? NULL : &consumed,
+                                               &self->cd);
+    return codec_tuple(decoded, consumed);
+}
+
+static int
+iconv_decoder_traverse(PyObject *op, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(op));
+    return 0;
+}
+
+static void
+iconv_decoder_dealloc(PyObject *op)
+{
+    iconv_decoder_object *self = iconv_decoder_CAST(op);
+    PyTypeObject *tp = Py_TYPE(self);
+    PyObject_GC_UnTrack(self);
+    if (self->cd != (iconv_t)-1) {
+        iconv_close(self->cd);
+    }
+    PyMem_Free(self->encoding);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyMethodDef iconv_decoder_methods[] = {
+    _CODECS_ICONVDECODER_DECODE_METHODDEF
+    {NULL, NULL}
+};
+
+static PyType_Slot iconv_decoder_slots[] = {
+    {Py_tp_new, _codecs_IconvDecoder},
+    {Py_tp_dealloc, iconv_decoder_dealloc},
+    {Py_tp_traverse, iconv_decoder_traverse},
+    {Py_tp_methods, iconv_decoder_methods},
+    {Py_tp_doc, (void *)_codecs_IconvDecoder__doc__},
+    {0, NULL}
+};
+
+static PyType_Spec iconv_decoder_spec = {
+    .name = "_codecs.IconvDecoder",
+    .basicsize = sizeof(iconv_decoder_object),
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE
+              | Py_TPFLAGS_HAVE_GC),
+    .slots = iconv_decoder_slots,
+};
+
 /*[clinic input]
 _codecs.iconv_decode
     encoding: str
@@ -661,7 +792,7 @@ _codecs_iconv_decode_impl(PyObject *module, const char *encoding,
     Py_ssize_t consumed = data->len;
     PyObject *decoded = _PyUnicode_DecodeIconv(encoding, data->buf, data->len,
                                                errors,
-                                               final ? NULL : &consumed);
+                                               final ? NULL : &consumed, NULL);
     return codec_tuple(decoded, consumed);
 }
 
@@ -1162,23 +1293,60 @@ static PyMethodDef _codecs_functions[] = {
     {NULL, NULL}                /* sentinel */
 };
 
+static int
+_codecs_exec(PyObject *module)
+{
+#ifdef HAVE_ICONV
+    _codecs_state *state = get_codecs_state(module);
+    state->IconvDecoderType = (PyTypeObject *)PyType_FromModuleAndSpec(
+            module, &iconv_decoder_spec, NULL);
+    if (state->IconvDecoderType == NULL) {
+        return -1;
+    }
+    if (PyModule_AddType(module, state->IconvDecoderType) < 0) {
+        return -1;
+    }
+#endif
+    return 0;
+}
+
+static int
+_codecs_traverse(PyObject *module, visitproc visit, void *arg)
+{
+    Py_VISIT(get_codecs_state(module)->IconvDecoderType);
+    return 0;
+}
+
+static int
+_codecs_clear(PyObject *module)
+{
+    Py_CLEAR(get_codecs_state(module)->IconvDecoderType);
+    return 0;
+}
+
+static void
+_codecs_free(void *module)
+{
+    (void)_codecs_clear((PyObject *)module);
+}
+
 static PyModuleDef_Slot _codecs_slots[] = {
     _Py_ABI_SLOT,
+    {Py_mod_exec, _codecs_exec},
     {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
     {Py_mod_gil, Py_MOD_GIL_NOT_USED},
     {0, NULL}
 };
 
 static struct PyModuleDef codecsmodule = {
-        PyModuleDef_HEAD_INIT,
-        "_codecs",
-        NULL,
-        0,
-        _codecs_functions,
-        _codecs_slots,
-        NULL,
-        NULL,
-        NULL
+        .m_base = PyModuleDef_HEAD_INIT,
+        .m_name = "_codecs",
+        .m_size = sizeof(_codecs_state),
+        .m_methods = _codecs_functions,
+        .m_slots = _codecs_slots,
+        .m_traverse = _codecs_traverse,
+        .m_clear = _codecs_clear,
+        .m_free = _codecs_free,
 };
 
 PyMODINIT_FUNC