Add PyUnicodeWriter_WriteString()

Author: vstinner

Include/cpython/unicodeobject.h

@@ -458,10 +458,10 @@ PyAPI_FUNC(void) PyUnicodeWriter_SetOverallocate(
 PyAPI_FUNC(int) PyUnicodeWriter_WriteChar(
     PyUnicodeWriter *writer,
     Py_UCS4 ch);
-PyAPI_FUNC(int) PyUnicodeWriter_WriteASCIIString(
+PyAPI_FUNC(int) PyUnicodeWriter_WriteString(
     PyUnicodeWriter *writer,
-    const char *ascii,
-    Py_ssize_t len);
+    const char *str,
+    Py_ssize_t size);
 PyAPI_FUNC(int) PyUnicodeWriter_WriteStr(
     PyUnicodeWriter *writer,
     PyObject *str);

Include/internal/pycore_unicodeobject.h

@@ -345,6 +345,14 @@ struct _PyUnicodeWriter {
     unsigned char readonly;
 };
 
+static inline void*
+_PyUnicodeWriter_GetDataEnd(_PyUnicodeWriter *writer)
+{
+    char *data = writer->data;
+    data += writer->pos * writer->kind;
+    return data;
+}
+
 // Initialize a Unicode writer.
 //
 // By default, the minimum buffer size is 0 character and overallocation is

Modules/_testcapi/unicode.c

@@ -259,8 +259,8 @@ test_unicodewriter(PyObject *self, PyObject *Py_UNUSED(args))
         goto error;
     }
 
-    // test PyUnicodeWriter_WriteASCIIString()
-    if (PyUnicodeWriter_WriteASCIIString(writer, " value", 6) < 0) {
+    // test PyUnicodeWriter_WriteString()
+    if (PyUnicodeWriter_WriteString(writer, " value", 6) < 0) {
         goto error;
     }
 
@@ -279,6 +279,65 @@ test_unicodewriter(PyObject *self, PyObject *Py_UNUSED(args))
 }
 
 
+static PyObject *
+test_unicodewriter_utf8(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    if (writer == NULL) {
+        return NULL;
+    }
+    if (PyUnicodeWriter_WriteString(writer, "ascii", -1) < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteChar(writer, '-') < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteString(writer, "latin1=\xC3\xA9", -1) < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteChar(writer, '-') < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteString(writer, "euro=\xE2\x82\xAC", -1) < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteChar(writer, '.') < 0) {
+        goto error;
+    }
+
+    PyObject *result = PyUnicodeWriter_Finish(writer);
+    if (result == NULL) {
+        return NULL;
+    }
+    assert(PyUnicode_EqualToUTF8(result,
+                                 "ascii-latin1=\xC3\xA9-euro=\xE2\x82\xAC."));
+    Py_DECREF(result);
+
+    Py_RETURN_NONE;
+
+error:
+    PyUnicodeWriter_Free(writer);
+    return NULL;
+}
+
+
+static PyObject *
+test_unicodewriter_invalid_utf8(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    if (writer == NULL) {
+        return NULL;
+    }
+    assert(PyUnicodeWriter_WriteString(writer, "invalid=\xFF", -1) < 0);
+    PyUnicodeWriter_Free(writer);
+
+    assert(PyErr_ExceptionMatches(PyExc_UnicodeDecodeError));
+    PyErr_Clear();
+
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *
 test_unicodewriter_format(PyObject *self, PyObject *Py_UNUSED(args))
 {
@@ -321,6 +380,8 @@ static PyMethodDef TestMethods[] = {
     {"unicode_asutf8",           unicode_asutf8,                 METH_VARARGS},
     {"unicode_copycharacters",   unicode_copycharacters,         METH_VARARGS},
     {"test_unicodewriter",       test_unicodewriter,             METH_NOARGS},
+    {"test_unicodewriter_utf8",  test_unicodewriter_utf8,        METH_NOARGS},
+    {"test_unicodewriter_invalid_utf8", test_unicodewriter_invalid_utf8, METH_NOARGS},
     {"test_unicodewriter_format", test_unicodewriter_format,     METH_NOARGS},
     {NULL},
 };

Objects/unicodeobject.c

@@ -4721,8 +4721,9 @@ ascii_decode(const char *start, const char *end, Py_UCS1 *dest)
     const char *p = start;
 
 #if SIZEOF_SIZE_T <= SIZEOF_VOID_P
-    assert(_Py_IS_ALIGNED(dest, ALIGNOF_SIZE_T));
-    if (_Py_IS_ALIGNED(p, ALIGNOF_SIZE_T)) {
+    if (_Py_IS_ALIGNED(p, ALIGNOF_SIZE_T)
+        && _Py_IS_ALIGNED(dest, ALIGNOF_SIZE_T))
+    {
         /* Fast path, see in STRINGLIB(utf8_decode) for
            an explanation. */
         /* Help allocation */
@@ -4945,6 +4946,7 @@ unicode_decode_utf8(const char *s, Py_ssize_t size,
 }
 
 
+// Used by PyUnicodeWriter_WriteString() implementation
 static int
 unicode_decode_utf8_writer(_PyUnicodeWriter *writer,
                            const char *s, Py_ssize_t size,
@@ -13525,12 +13527,15 @@ _PyUnicodeWriter_WriteASCIIString(_PyUnicodeWriter *writer,
 }
 
 int
-PyUnicodeWriter_WriteASCIIString(PyUnicodeWriter *writer,
-                                 const char *ascii,
-                                 Py_ssize_t len)
+PyUnicodeWriter_WriteString(PyUnicodeWriter *writer,
+                            const char *str,
+                            Py_ssize_t size)
 {
-    return _PyUnicodeWriter_WriteASCIIString((_PyUnicodeWriter*)writer,
-                                             ascii, len);
+    if (size == -1) {
+        size = strlen(str);
+    }
+    return unicode_decode_utf8_writer((_PyUnicodeWriter*)writer, str, size,
+                                      _Py_ERROR_STRICT, NULL, NULL);
 }
 
 int