PyUnicodeWriter_Create() expects a length

vstinner · vstinner · commit b12f085fe112 · 2024-06-07T21:29:20.000+02:00
Remove PyUnicodeWriter_SetOverallocate().
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)

Original file line number	Diff line number	Diff line change
`@@ -224,14 +224,11 @@ unicode_copycharacters(PyObject self, PyObject args)`
`224`	`224`	`static PyObject *`
`225`	`225`	`test_unicodewriter(PyObject self, PyObject Py_UNUSED(args))`
`226`	`226`	`{`
`227`		`- PyUnicodeWriter *writer = PyUnicodeWriter_Create();`
	`227`	`+ PyUnicodeWriter *writer = PyUnicodeWriter_Create(100);`
`228`	`228`	`if (writer == NULL) {`
`229`	`229`	`return NULL;`
`230`	`230`	`}`
`231`	`231`
`232`		`- // test PyUnicodeWriter_SetOverallocate()`
`233`		`- PyUnicodeWriter_SetOverallocate(writer, 1);`
`234`		`-`
`235`	`232`	`// test PyUnicodeWriter_WriteUTF8()`
`236`	`233`	`if (PyUnicodeWriter_WriteUTF8(writer, "var", -1) < 0) {`
`237`	`234`	`goto error;`
`@@ -293,7 +290,7 @@ test_unicodewriter(PyObject self, PyObject Py_UNUSED(args))`
`293`	`290`	`static PyObject *`
`294`	`291`	`test_unicodewriter_utf8(PyObject self, PyObject Py_UNUSED(args))`
`295`	`292`	`{`
`296`		`- PyUnicodeWriter *writer = PyUnicodeWriter_Create();`
	`293`	`+ PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);`
`297`	`294`	`if (writer == NULL) {`
`298`	`295`	`return NULL;`
`299`	`296`	`}`
`@@ -335,7 +332,7 @@ test_unicodewriter_utf8(PyObject self, PyObject Py_UNUSED(args))`
`335`	`332`	`static PyObject *`
`336`	`333`	`test_unicodewriter_invalid_utf8(PyObject self, PyObject Py_UNUSED(args))`
`337`	`334`	`{`
`338`		`- PyUnicodeWriter *writer = PyUnicodeWriter_Create();`
	`335`	`+ PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);`
`339`	`336`	`if (writer == NULL) {`
`340`	`337`	`return NULL;`
`341`	`338`	`}`
`@@ -352,7 +349,7 @@ test_unicodewriter_invalid_utf8(PyObject self, PyObject Py_UNUSED(args))`
`352`	`349`	`static PyObject *`
`353`	`350`	`test_unicodewriter_format(PyObject self, PyObject Py_UNUSED(args))`
`354`	`351`	`{`
`355`		`- PyUnicodeWriter *writer = PyUnicodeWriter_Create();`
	`352`	`+ PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);`
`356`	`353`	`if (writer == NULL) {`
`357`	`354`	`return NULL;`
`358`	`355`	`}`