gh-119182: Add PyUnicodeWriter_WriteWideChar()

Author: vstinner

Include/cpython/unicodeobject.h

@@ -459,6 +459,10 @@ PyAPI_FUNC(int) PyUnicodeWriter_WriteUTF8(
     PyUnicodeWriter *writer,
     const char *str,
     Py_ssize_t size);
+PyAPI_FUNC(int) PyUnicodeWriter_WriteWideChar(
+    PyUnicodeWriter *writer,
+    wchar_t *str,
+    Py_ssize_t size);
 
 PyAPI_FUNC(int) PyUnicodeWriter_WriteStr(
     PyUnicodeWriter *writer,

Modules/_testcapi/unicode.c

@@ -436,6 +436,42 @@ test_unicodewriter_format_recover_error(PyObject *self, PyObject *Py_UNUSED(args
 }
 
 
+static PyObject *
+test_unicodewriter_widechar(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
+    if (writer == NULL) {
+        return NULL;
+    }
+    if (PyUnicodeWriter_WriteWideChar(writer, L"latin1=\xE9 IGNORED", 8) < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteWideChar(writer, L"-", 1) < 0) {
+        goto error;
+    }
+    if (PyUnicodeWriter_WriteWideChar(writer, L"euro=\u20AC", -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,
+                                 "latin1=\xC3\xA9-euro=\xE2\x82\xAC."));
+    Py_DECREF(result);
+
+    Py_RETURN_NONE;
+
+error:
+    PyUnicodeWriter_Discard(writer);
+    return NULL;
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"unicode_new",              unicode_new,                    METH_VARARGS},
     {"unicode_fill",             unicode_fill,                   METH_VARARGS},
@@ -450,6 +486,7 @@ static PyMethodDef TestMethods[] = {
     {"test_unicodewriter_recover_error", test_unicodewriter_recover_error, METH_NOARGS},
     {"test_unicodewriter_format", test_unicodewriter_format,     METH_NOARGS},
     {"test_unicodewriter_format_recover_error", test_unicodewriter_format_recover_error, METH_NOARGS},
+    {"test_unicodewriter_widechar", test_unicodewriter_widechar, METH_NOARGS},
     {NULL},
 };
 

Objects/unicodeobject.c

@@ -13500,6 +13500,27 @@ PyUnicodeWriter_WriteUTF8(PyUnicodeWriter *writer,
     return res;
 }
 
+
+int
+PyUnicodeWriter_WriteWideChar(PyUnicodeWriter *writer,
+                              wchar_t *str,
+                              Py_ssize_t size)
+{
+    if (size < 0) {
+        size = wcslen(str);
+    }
+    PyObject *obj = PyUnicode_FromWideChar(str, size);
+    if (obj == NULL) {
+        return -1;
+    }
+
+    _PyUnicodeWriter *_writer = (_PyUnicodeWriter *)writer;
+    int res = _PyUnicodeWriter_WriteStr(_writer, obj);
+    Py_DECREF(obj);
+    return res;
+}
+
+
 int
 _PyUnicodeWriter_WriteLatin1String(_PyUnicodeWriter *writer,
                                    const char *str, Py_ssize_t len)