WIP: Add PyUnicodeWriter API

TODO: update API doc.

Author: vstinner

docs/changelog.rst

@@ -1,6 +1,18 @@
 Changelog
 =========
 
+* 2024-05-23: Add functions:
+
+  * ``PyUnicodeWriter_Create()``
+  * ``PyUnicodeWriter_Free()``
+  * ``PyUnicodeWriter_Finish()``
+  * ``PyUnicodeWriter_SetOverallocate()``
+  * ``PyUnicodeWriter_WriteChar()``
+  * ``PyUnicodeWriter_WriteUTF8()``
+  * ``PyUnicodeWriter_WriteString()``
+  * ``PyUnicodeWriter_WriteSubstring()``
+  * ``PyUnicodeWriter_Format()``
+
 * 2024-04-02: Add ``PyDict_SetDefaultRef()`` function.
 * 2024-03-29: Add ``PyList_GetItemRef()`` function.
 * 2024-03-21: Add functions:

pythoncapi_compat.h

@@ -1338,6 +1338,112 @@ PyDict_SetDefaultRef(PyObject *d, PyObject *key, PyObject *default_value,
 }
 #endif
 
+#if PY_VERSION_HEX < 0x030E0000
+typedef struct PyUnicodeWriter PyUnicodeWriter;
+
+static inline PyUnicodeWriter* PyUnicodeWriter_Create(void)
+{
+    const size_t size = sizeof(_PyUnicodeWriter);
+    PyUnicodeWriter *writer = (PyUnicodeWriter *)PyMem_Malloc(size);
+    if (writer == _Py_NULL) {
+        PyErr_NoMemory();
+        return _Py_NULL;
+    }
+    _PyUnicodeWriter_Init((_PyUnicodeWriter*)writer);
+    ((_PyUnicodeWriter*)writer)->overallocate = 1;
+    return writer;
+}
+
+static inline void PyUnicodeWriter_Free(PyUnicodeWriter *writer)
+{
+    _PyUnicodeWriter_Dealloc((_PyUnicodeWriter*)writer);
+    PyMem_Free(writer);
+}
+
+static inline PyObject* PyUnicodeWriter_Finish(PyUnicodeWriter *writer)
+{
+    PyObject *str = _PyUnicodeWriter_Finish((_PyUnicodeWriter*)writer);
+    PyMem_Free(writer);
+    return str;
+}
+
+static inline void
+PyUnicodeWriter_SetOverallocate(PyUnicodeWriter *writer, int overallocate)
+{
+    ((_PyUnicodeWriter*)writer)->overallocate = (unsigned char)overallocate;
+}
+
+static inline int
+PyUnicodeWriter_WriteChar(PyUnicodeWriter *writer, Py_UCS4 ch)
+{
+    return _PyUnicodeWriter_WriteChar((_PyUnicodeWriter*)writer, ch);
+}
+
+static inline int
+PyUnicodeWriter_WriteString(PyUnicodeWriter *writer, PyObject *str)
+{
+    if (!PyUnicode_Check(str)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expect str, not %s", Py_TYPE(str)->tp_name);
+        return -1;
+    }
+    return _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str);
+}
+
+static inline int
+PyUnicodeWriter_WriteUTF8(PyUnicodeWriter *writer,
+                          const char *str, Py_ssize_t size)
+{
+    if (size < 0) {
+        size = (Py_ssize_t)strlen(str);
+    }
+    PyObject *str_obj = PyUnicode_FromStringAndSize(str, size);
+    if (str_obj == _Py_NULL) {
+        return -1;
+    }
+
+    int res = PyUnicodeWriter_WriteString(writer, str_obj);
+    Py_DECREF(str_obj);
+    return res;
+}
+
+static inline int
+PyUnicodeWriter_WriteSubstring(PyUnicodeWriter *writer, PyObject *str,
+                               Py_ssize_t start, Py_ssize_t end)
+{
+    if (!PyUnicode_Check(str)) {
+        PyErr_Format(PyExc_TypeError, "expect str, not %T", str);
+        return -1;
+    }
+    if (start < 0 || start > end) {
+        PyErr_Format(PyExc_ValueError, "invalid start argument");
+        return -1;
+    }
+    if (end > PyUnicode_GET_LENGTH(str)) {
+        PyErr_Format(PyExc_ValueError, "invalid end argument");
+        return -1;
+    }
+
+    return _PyUnicodeWriter_WriteSubstring((_PyUnicodeWriter*)writer, str,
+                                           start, end);
+}
+
+static inline int
+PyUnicodeWriter_Format(PyUnicodeWriter *writer, const char *format, ...)
+{
+    va_list vargs;
+    va_start(vargs, format);
+    PyObject *str = PyUnicode_FromFormatV(format, vargs);
+    va_end(vargs);
+    if (str == _Py_NULL) {
+        return -1;
+    }
+
+    int res = PyUnicodeWriter_WriteString(writer, str);
+    Py_DECREF(str);
+    return res;
+}
+#endif  // PY_VERSION_HEX < 0x030E0000
 
 #ifdef __cplusplus
 }

tests/test_pythoncapi_compat_cext.c

@@ -1724,6 +1724,102 @@ test_get_constant(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+static PyObject *
+test_unicodewriter(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
+{
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    if (writer == NULL) {
+        return NULL;
+    }
+    int ret;
+
+    // test PyUnicodeWriter_SetOverallocate()
+    PyUnicodeWriter_SetOverallocate(writer, 1);
+
+    // test PyUnicodeWriter_WriteString()
+    PyObject *str = PyUnicode_FromString("var");
+    if (str == NULL) {
+        goto error;
+    }
+    ret = PyUnicodeWriter_WriteString(writer, str);
+    Py_CLEAR(str);
+    if (ret < 0) {
+        goto error;
+    }
+
+    // test PyUnicodeWriter_WriteChar()
+    if (PyUnicodeWriter_WriteChar(writer, '=') < 0) {
+        goto error;
+    }
+
+    // test PyUnicodeWriter_WriteSubstring()
+    str = PyUnicode_FromString("[long]");
+    if (str == NULL) {
+        goto error;
+    }
+    ret = PyUnicodeWriter_WriteSubstring(writer, str, 1, 5);
+    Py_CLEAR(str);
+    if (ret < 0) {
+        goto error;
+    }
+
+    // test PyUnicodeWriter_WriteUTF8()
+    if (PyUnicodeWriter_WriteUTF8(writer, " valu\xC3\xA9", -1) < 0) {
+        goto error;
+    }
+
+    {
+        PyObject *result = PyUnicodeWriter_Finish(writer);
+        if (result == NULL) {
+            return NULL;
+        }
+        assert(PyUnicode_EqualToUTF8(result, "var=long valu\xC3\xA9"));
+        Py_DECREF(result);
+    }
+
+    Py_RETURN_NONE;
+
+error:
+    PyUnicodeWriter_Free(writer);
+    return NULL;
+}
+
+
+static PyObject *
+test_unicodewriter_format(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
+{
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    if (writer == NULL) {
+        return NULL;
+    }
+
+    // test PyUnicodeWriter_Format()
+    if (PyUnicodeWriter_Format(writer, "%s %i", "Hello", 123) < 0) {
+        goto error;
+    }
+
+    // test PyUnicodeWriter_WriteChar()
+    if (PyUnicodeWriter_WriteChar(writer, '.') < 0) {
+        goto error;
+    }
+
+    {
+        PyObject *result = PyUnicodeWriter_Finish(writer);
+        if (result == NULL) {
+            return NULL;
+        }
+        assert(PyUnicode_EqualToUTF8(result, "Hello 123."));
+        Py_DECREF(result);
+    }
+
+    Py_RETURN_NONE;
+
+error:
+    PyUnicodeWriter_Free(writer);
+    return NULL;
+}
+
+
 static struct PyMethodDef methods[] = {
     {"test_object", test_object, METH_NOARGS, _Py_NULL},
     {"test_py_is", test_py_is, METH_NOARGS, _Py_NULL},
@@ -1762,6 +1858,8 @@ static struct PyMethodDef methods[] = {
     {"test_time", test_time, METH_NOARGS, _Py_NULL},
 #endif
     {"test_get_constant", test_get_constant, METH_NOARGS, _Py_NULL},
+    {"test_unicodewriter", test_unicodewriter, METH_NOARGS, _Py_NULL},
+    {"test_unicodewriter_format", test_unicodewriter_format, METH_NOARGS, _Py_NULL},
     {_Py_NULL, _Py_NULL, 0, _Py_NULL}
 };