Issue #20404: reject non-text encodings early in TextIOWrapper.

1 files changed, 63 insertions, 21 deletions

diff --git a/Python/codecs.c b/Python/codecs.c
index 5ebc4cb..0b736c1 100644
--- a/Python/codecs.c
+++ b/Python/codecs.c
@@ -243,20 +243,15 @@ PyObject *codec_getitem(const char *encoding, int index)
     return v;
 }
 
-/* Helper function to create an incremental codec. */
-
+/* Helper functions to create an incremental codec. */
 static
-PyObject *codec_getincrementalcodec(const char *encoding,
-                                    const char *errors,
-                                    const char *attrname)
+PyObject *codec_makeincrementalcodec(PyObject *codec_info,
+                                     const char *errors,
+                                     const char *attrname)
 {
-    PyObject *codecs, *ret, *inccodec;
+    PyObject *ret, *inccodec;
 
-    codecs = _PyCodec_Lookup(encoding);
-    if (codecs == NULL)
-        return NULL;
-    inccodec = PyObject_GetAttrString(codecs, attrname);
-    Py_DECREF(codecs);
+    inccodec = PyObject_GetAttrString(codec_info, attrname);
     if (inccodec == NULL)
         return NULL;
     if (errors)
@@ -267,6 +262,21 @@ PyObject *codec_getincrementalcodec(const char *encoding,
     return ret;
 }
 
+static
+PyObject *codec_getincrementalcodec(const char *encoding,
+                                    const char *errors,
+                                    const char *attrname)
+{
+    PyObject *codec_info, *ret;
+
+    codec_info = _PyCodec_Lookup(encoding);
+    if (codec_info == NULL)
+        return NULL;
+    ret = codec_makeincrementalcodec(codec_info, errors, attrname);
+    Py_DECREF(codec_info);
+    return ret;
+}
+
 /* Helper function to create a stream codec. */
 
 static
@@ -290,6 +300,24 @@ PyObject *codec_getstreamcodec(const char *encoding,
     return streamcodec;
 }
 
+/* Helpers to work with the result of _PyCodec_Lookup
+
+ */
+PyObject *_PyCodecInfo_GetIncrementalDecoder(PyObject *codec_info,
+                                             const char *errors)
+{
+    return codec_makeincrementalcodec(codec_info, errors,
+                                      "incrementaldecoder");
+}
+
+PyObject *_PyCodecInfo_GetIncrementalEncoder(PyObject *codec_info,
+                                             const char *errors)
+{
+    return codec_makeincrementalcodec(codec_info, errors,
+                                      "incrementalencoder");
+}
+
+
 /* Convenience APIs to query the Codec registry.
 
    All APIs return a codec object with incremented refcount.
@@ -447,15 +475,12 @@ PyObject *PyCodec_Decode(PyObject *object,
 }
 
 /* Text encoding/decoding API */
-static
-PyObject *codec_getitem_checked(const char *encoding,
-                                const char *operation_name,
-                                int index)
+PyObject * _PyCodec_LookupTextEncoding(const char *encoding,
+                                       const char *alternate_command)
 {
     _Py_IDENTIFIER(_is_text_encoding);
     PyObject *codec;
     PyObject *attr;
-    PyObject *v;
     int is_text_codec;
 
     codec = _PyCodec_Lookup(encoding);
@@ -482,27 +507,44 @@ PyObject *codec_getitem_checked(const char *encoding,
                 Py_DECREF(codec);
                 PyErr_Format(PyExc_LookupError,
                              "'%.400s' is not a text encoding; "
-                             "use codecs.%s() to handle arbitrary codecs",
-                             encoding, operation_name);
+                             "use %s to handle arbitrary codecs",
+                             encoding, alternate_command);
                 return NULL;
             }
         }
     }
 
+    /* This appears to be a valid text encoding */
+    return codec;
+}
+
+
+static
+PyObject *codec_getitem_checked(const char *encoding,
+                                const char *alternate_command,
+                                int index)
+{
+    PyObject *codec;
+    PyObject *v;
+
+    codec = _PyCodec_LookupTextEncoding(encoding, alternate_command);
+    if (codec == NULL)
+        return NULL;
+
     v = PyTuple_GET_ITEM(codec, index);
-    Py_DECREF(codec);
     Py_INCREF(v);
+    Py_DECREF(codec);
     return v;
 }
 
 static PyObject * _PyCodec_TextEncoder(const char *encoding)
 {
-    return codec_getitem_checked(encoding, "encode", 0);
+    return codec_getitem_checked(encoding, "codecs.encode()", 0);
 }
 
 static PyObject * _PyCodec_TextDecoder(const char *encoding)
 {
-    return codec_getitem_checked(encoding, "decode", 1);
+    return codec_getitem_checked(encoding, "codecs.decode()", 1);
 }
 
 PyObject *_PyCodec_EncodeText(PyObject *object,