PyUnicodeWriter_Create() expects a length

Remove PyUnicodeWriter_SetOverallocate().
vstinner · Jun 5, 2024 · fd7432e · fd7432e
1 parent acf42e3
commit fd7432e
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 26 deletions.
diff --git a/Include/cpython/unicodeobject.h b/Include/cpython/unicodeobject.h
@@ -448,14 +448,10 @@ PyAPI_FUNC(PyObject*) PyUnicode_FromKindAndData(
 
 typedef struct PyUnicodeWriter PyUnicodeWriter;
 
-PyAPI_FUNC(PyUnicodeWriter*) PyUnicodeWriter_Create(void);
+PyAPI_FUNC(PyUnicodeWriter*) PyUnicodeWriter_Create(Py_ssize_t length);
 PyAPI_FUNC(void) PyUnicodeWriter_Discard(PyUnicodeWriter *writer);
 PyAPI_FUNC(PyObject*) PyUnicodeWriter_Finish(PyUnicodeWriter *writer);
 
-PyAPI_FUNC(void) PyUnicodeWriter_SetOverallocate(
-    PyUnicodeWriter *writer,
-    int overallocate);
-
 PyAPI_FUNC(int) PyUnicodeWriter_WriteChar(
     PyUnicodeWriter *writer,
     Py_UCS4 ch);

diff --git a/Modules/_testcapi/unicode.c b/Modules/_testcapi/unicode.c
@@ -224,14 +224,11 @@ unicode_copycharacters(PyObject *self, PyObject *args)
 static PyObject *
 test_unicodewriter(PyObject *self, PyObject *Py_UNUSED(args))
 {
-    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(100);
     if (writer == NULL) {
         return NULL;
     }
 
-    // test PyUnicodeWriter_SetOverallocate()
-    PyUnicodeWriter_SetOverallocate(writer, 1);
-
     // test PyUnicodeWriter_WriteUTF8()
     if (PyUnicodeWriter_WriteUTF8(writer, "var", -1) < 0) {
         goto error;
@@ -293,7 +290,7 @@ test_unicodewriter(PyObject *self, PyObject *Py_UNUSED(args))
 static PyObject *
 test_unicodewriter_utf8(PyObject *self, PyObject *Py_UNUSED(args))
 {
-    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
     if (writer == NULL) {
         return NULL;
     }
@@ -335,7 +332,7 @@ test_unicodewriter_utf8(PyObject *self, PyObject *Py_UNUSED(args))
 static PyObject *
 test_unicodewriter_invalid_utf8(PyObject *self, PyObject *Py_UNUSED(args))
 {
-    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
     if (writer == NULL) {
         return NULL;
     }
@@ -352,7 +349,7 @@ test_unicodewriter_invalid_utf8(PyObject *self, PyObject *Py_UNUSED(args))
 static PyObject *
 test_unicodewriter_format(PyObject *self, PyObject *Py_UNUSED(args))
 {
-    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
     if (writer == NULL) {
         return NULL;
     }

diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
@@ -13139,22 +13139,29 @@ _PyUnicodeWriter_Init(_PyUnicodeWriter *writer)
     /* use a kind value smaller than PyUnicode_1BYTE_KIND so
        _PyUnicodeWriter_PrepareKind() will copy the buffer. */
     assert(writer->kind == 0);
-    assert(writer->kind <= PyUnicode_1BYTE_KIND);
+    assert(writer->kind < PyUnicode_1BYTE_KIND);
 }
 
 
 PyUnicodeWriter*
-PyUnicodeWriter_Create(void)
+PyUnicodeWriter_Create(Py_ssize_t length)
 {
     const size_t size = sizeof(_PyUnicodeWriter);
-    PyUnicodeWriter *writer = (PyUnicodeWriter *)PyMem_Malloc(size);
-    if (writer == NULL) {
+    PyUnicodeWriter *pub_writer = (PyUnicodeWriter *)PyMem_Malloc(size);
+    if (pub_writer == NULL) {
         PyErr_NoMemory();
         return NULL;
     }
-    _PyUnicodeWriter_Init((_PyUnicodeWriter*)writer);
-    PyUnicodeWriter_SetOverallocate(writer, 1);
-    return writer;
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter *)pub_writer;
+
+    _PyUnicodeWriter_Init(writer);
+    if (_PyUnicodeWriter_Prepare(writer, length, 127) < 0) {
+        PyUnicodeWriter_Discard(pub_writer);
+        return NULL;
+    }
+    writer->overallocate = 1;
+
+    return pub_writer;
 }
 
 
@@ -13176,13 +13183,6 @@ _PyUnicodeWriter_InitWithBuffer(_PyUnicodeWriter *writer, PyObject *buffer)
 }
 
 
-void
-PyUnicodeWriter_SetOverallocate(PyUnicodeWriter *writer, int overallocate)
-{
-    ((_PyUnicodeWriter*)writer)->overallocate = (unsigned char)overallocate;
-}
-
-
 int
 _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
                                  Py_ssize_t length, Py_UCS4 maxchar)