Issue #16145: Support legacy strings in the _csv module.

3 files changed, 50 insertions, 6 deletions

diff --git a/Lib/test/test_csv.py b/Lib/test/test_csv.py
index 55796a2..96f8aa7 100644
--- a/Lib/test/test_csv.py
+++ b/Lib/test/test_csv.py
@@ -197,6 +197,17 @@ class Test_Csv(unittest.TestCase):
             fileobj.seek(0)
             self.assertEqual(fileobj.read(), "a,b\r\nc,d\r\n")
 
+    @support.cpython_only
+    def test_writerows_legacy_strings(self):
+        import _testcapi
+
+        c = _testcapi.unicode_legacy_string('a')
+        with TemporaryFile("w+", newline='') as fileobj:
+            writer = csv.writer(fileobj)
+            writer.writerows([[c]])
+            fileobj.seek(0)
+            self.assertEqual(fileobj.read(), "a\r\n")
+
     def _read_test(self, input, expect, **kwargs):
         reader = csv.reader(input, **kwargs)
         result = list(reader)
diff --git a/Modules/_csv.c b/Modules/_csv.c
index cc87bad..48a5cf8 100644
--- a/Modules/_csv.c
+++ b/Modules/_csv.c
@@ -13,8 +13,6 @@ module instead.
 #include "Python.h"
 #include "structmember.h"
 
-#define IS_BASESTRING(o) \
-    PyUnicode_Check(o)
 
 typedef struct {
     PyObject *error_obj;   /* CSV exception */
@@ -248,6 +246,7 @@ _set_char(const char *name, Py_UCS4 *target, PyObject *src, Py_UCS4 dflt)
                     name);
                 return -1;
             }
+            /* PyUnicode_READY() is called in PyUnicode_GetLength() */
             if (len > 0)
                 *target = PyUnicode_READ_CHAR(src, 0);
         }
@@ -263,12 +262,14 @@ _set_str(const char *name, PyObject **target, PyObject *src, const char *dflt)
     else {
         if (src == Py_None)
             *target = NULL;
-        else if (!IS_BASESTRING(src)) {
+        else if (!PyUnicode_Check(src)) {
             PyErr_Format(PyExc_TypeError,
                          "\"%s\" must be a string", name);
             return -1;
         }
         else {
+            if (PyUnicode_READY(src) == -1)
+                return -1;
             Py_XDECREF(*target);
             Py_INCREF(src);
             *target = src;
@@ -357,7 +358,7 @@ dialect_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
         return NULL;
 
     if (dialect != NULL) {
-        if (IS_BASESTRING(dialect)) {
+        if (PyUnicode_Check(dialect)) {
             dialect = get_dialect_from_registry(dialect);
             if (dialect == NULL)
                 return NULL;
@@ -808,6 +809,10 @@ Reader_iternext(ReaderObj *self)
             Py_DECREF(lineobj);
             return NULL;
         }
+        if (PyUnicode_READY(lineobj) == -1) {
+            Py_DECREF(lineobj);
+            return NULL;
+        }
         ++self->line_num;
         kind = PyUnicode_KIND(lineobj);
         data = PyUnicode_DATA(lineobj);
@@ -1108,6 +1113,8 @@ join_append(WriterObj *self, PyObject *field, int *quoted, int quote_empty)
     Py_ssize_t rec_len;
 
     if (field != NULL) {
+        if (PyUnicode_READY(field) == -1)
+            return 0;
         field_kind = PyUnicode_KIND(field);
         field_data = PyUnicode_DATA(field);
         field_len = PyUnicode_GET_LENGTH(field);
@@ -1403,11 +1410,13 @@ csv_register_dialect(PyObject *module, PyObject *args, PyObject *kwargs)
 
     if (!PyArg_UnpackTuple(args, "", 1, 2, &name_obj, &dialect_obj))
         return NULL;
-    if (!IS_BASESTRING(name_obj)) {
+    if (!PyUnicode_Check(name_obj)) {
         PyErr_SetString(PyExc_TypeError,
-                        "dialect name must be a string or unicode");
+                        "dialect name must be a string");
         return NULL;
     }
+    if (PyUnicode_READY(name_obj) == -1)
+        return NULL;
     dialect = _call_dialect(dialect_obj, kwargs);
     if (dialect == NULL)
         return NULL;
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index ab11f51..2f43813 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -1521,6 +1521,29 @@ unicode_transformdecimaltoascii(PyObject *self, PyObject *args)
 }
 
 static PyObject *
+unicode_legacy_string(PyObject *self, PyObject *args)
+{
+    Py_UNICODE *data;
+    Py_ssize_t len;
+    PyObject *u;
+
+    if (!PyArg_ParseTuple(args, "u#", &data, &len))
+        return NULL;
+
+    u = PyUnicode_FromUnicode(NULL, len);
+    if (u == NULL)
+        return NULL;
+
+    memcpy(PyUnicode_AS_UNICODE(u), data, len * sizeof(Py_UNICODE));
+
+    if (len > 0) { /* The empty string is always ready. */
+        assert(!PyUnicode_IS_READY(u));
+    }
+
+    return u;
+}
+
+static PyObject *
 getargs_w_star(PyObject *self, PyObject *args)
 {
     Py_buffer buffer;
@@ -2506,6 +2529,7 @@ static PyMethodDef TestMethods[] = {
     {"unicode_aswidecharstring",unicode_aswidecharstring,        METH_VARARGS},
     {"unicode_encodedecimal",   unicode_encodedecimal,           METH_VARARGS},
     {"unicode_transformdecimaltoascii", unicode_transformdecimaltoascii, METH_VARARGS},
+    {"unicode_legacy_string",   unicode_legacy_string,           METH_VARARGS},
 #ifdef WITH_THREAD
     {"_test_thread_state",      test_thread_state,               METH_VARARGS},
     {"_pending_threadfunc",     pending_threadfunc,              METH_VARARGS},