gh-158451: Add _PyUnicodeWriter_FinishWithSize() test

Include/internal/pycore_unicodeobject.h

@@ -189,7 +189,8 @@ _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch)
     return 0;
 }
 
-extern PyObject* _PyUnicodeWriter_FinishWithSize(
+// Export for '_testcapi' shared extension
+PyAPI_FUNC(PyObject*) _PyUnicodeWriter_FinishWithSize(
     _PyUnicodeWriter *writer,
     Py_ssize_t size);
 

Lib/test/test_capi/test_unicode.py

@@ -2322,6 +2322,19 @@ def test_readonly_optim_large_int(self):
                 self.assertEqual(writer.get_buffer(), expected)
                 self.assertEqual(writer.finish(), large_int_str)
 
+    def test_finish_with_size(self):
+        # Test _PyUnicodeWriter_FinishWithSize(). Truncate text requires
+        # to change the buffer kind.
+        text = 'a\xff\u20ac\U0010ffff'
+        expected = (0x10ffff, False)  # do not test the buffer size
+        for size in range(len(text) + 1):
+            writer = self.create_writer(0)
+            writer.write_utf8(text.encode(), -1)
+            self.assertEqual(writer.get_buffer()[1:], expected)
+            self.assertEqual(writer.finish_with_size(size), text[:size])
+
+        # CRASHES writer.finish_with_size(len(text) + 1)
+
 
 # Test PyUnicodeWriter_Format()
 @unittest.skipIf(ctypes is None, 'need ctypes')

Modules/_testcapi/unicode.c

@@ -795,6 +795,31 @@ writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args))
 }
 
 
+static PyObject*
+writer_finish_with_size(PyObject *self_raw, PyObject *args)
+{
+    PyAPI_FUNC(PyObject*) _PyUnicodeWriter_FinishWithSize(
+        _PyUnicodeWriter *writer,
+        Py_ssize_t size);
+
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    Py_ssize_t size;
+    if (!PyArg_ParseTuple(args, "n", &size)) {
+        return NULL;
+    }
+
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter*)self->writer;
+    PyObject *str = _PyUnicodeWriter_FinishWithSize(writer, size);
+    PyUnicodeWriter_Discard(self->writer);
+    self->writer = NULL;
+    return str;
+}
+
+
 static PyMethodDef writer_methods[] = {
     {"write_char", _PyCFunction_CAST(writer_write_char), METH_VARARGS},
     {"write_utf8", _PyCFunction_CAST(writer_write_utf8), METH_VARARGS},
@@ -809,6 +834,7 @@ static PyMethodDef writer_methods[] = {
     {"get_pointer", _PyCFunction_CAST(writer_get_pointer), METH_VARARGS},
     {"get_buffer", _PyCFunction_CAST(writer_get_buffer), METH_VARARGS},
     {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
+    {"finish_with_size", _PyCFunction_CAST(writer_finish_with_size), METH_VARARGS},
     {NULL,              NULL}           /* sentinel */
 };