gh-157710: Defer allocation in PyUnicodeWriter_Create() (#157969)

PyUnicodeWriter_Create(length) no longer allocates 'length'
characters immediately. The allocation of the buffer is now done
lazily at the first write, except if the read-only optimization is
used. So the read-only optimization can also be used even if length
is greater than 0.

No longer overallocate the first buffer (at the first write).
Only overallocate when the buffer is resized (at the second write).
It avoids the need to truncate in PyUnicodeWriter_Finish() when
PyUnicodeWriter_Create(length) used the exact output length.

Add get_buffer() method to writer tests.

Author: vstinner

Lib/test/test_capi/test_unicode.py

@@ -1880,6 +1880,36 @@ def test_basic(self):
         self.assertEqual(writer.finish(),
                          "var=long value 'repr'")
 
+    def test_create(self):
+        # Test PyUnicodeWriter_Create() with non-zero size
+        s = 'Monty Python'
+
+        # Preallocate the exact length. Use 2 writes to force the creation
+        # of a buffer:
+        #   1. Use the read-only optimization.
+        #   2. Allocate a buffer of length character.
+        # No resize needed in finish().
+        writer = self.create_writer(len(s))
+        writer.write_str(s[:5])
+        self.assertEqual(writer.get_buffer(), (5, 127, True))
+        writer.write_str(s[5:])
+        self.assertEqual(writer.get_buffer(), (len(s), 127, False))
+        self.assertEqual(writer.finish(), s)
+
+        # Preallocate len(s)-1 characters. Use 3 writes:
+        #   1. Use read-only optimization.
+        #   2. Allocate a buffer of len-1 characters.
+        #   3. Resize the buffer with overallocation.
+        # finish() has to truncate the buffer.
+        writer = self.create_writer(len(s) - 1)
+        writer.write_str(s[:2])
+        self.assertEqual(writer.get_buffer(), (2, 127, True))
+        writer.write_str(s[2:5])
+        self.assertEqual(writer.get_buffer(), (len(s) - 1, 127, False))
+        writer.write_str(s[5:])
+        self.assertGreater(writer.get_buffer()[0], len(s))
+        self.assertEqual(writer.finish(), s)
+
     def test_repr_null(self):
         writer = self.create_writer(0)
         writer.write_utf8(b'var=', -1)
@@ -2087,32 +2117,39 @@ def test_substring_empty(self):
     def test_singletons(self):
         for size in (0, 123):
             with self.subTest(size=size):
+                # PyUnicodeWriter_Finish() returns the empty string singleton
+                # if no character has been written.
                 writer = self.create_writer(size)
                 writer.write_utf8(b'utf8', 0)
                 writer.write_ascii(b'ascii', 0)
                 writer.write_widechar(b'wstr', 0)
                 writer.write_ucs4(b'ucs4', 0)
-                writer.write_substring('text', 0, 0)
+                writer.write_substring('text', 2, 2)
+                self.assertEqual(writer.get_buffer(), (None, 127, False))
                 self.assertIs(writer.finish(), '')
 
         for size in (0, 123):
             for ch in range(256):
                 with self.subTest(size=size, ch=ch):
                     ch = chr(ch)
+                    maxchar = (255 if ord(ch) >= 128 else 127)
 
-                    # If the first write is a Latin1 character and no buffer
-                    # was allocated yet, use the singleton as the read-only
-                    # buffer
+                    # PyUnicodeWriter_WriteChar(ch) uses the read-only
+                    # optimization with the character singleton if ch is a
+                    # Latin1 character and no buffer was allocated yet.
                     writer = self.create_writer(size)
                     writer.write_char(ord(ch))
+                    self.assertEqual(writer.get_buffer(),
+                                     (1, maxchar, True))
                     self.assertIs(writer.finish(), ch)
 
-                    # PyUnicodeWriter_Finish() replaces the buffer
-                    # with the singleton
+                    # PyUnicodeWriter_Finish() replaces the buffer with the
+                    # singleton. Use PyUnicodeWriter_WriteSubstring() to avoid
+                    # the read-only buffer optimization.
                     writer = self.create_writer(size)
-                    # Use PyUnicodeWriter_WriteSubstring() to avoid
-                    # the read-only buffer optimization
-                    writer.write_substring(ch + 'xxx', 0, 1)
+                    writer.write_substring('xxx' + ch + 'y', 3, 4)
+                    self.assertEqual(writer.get_buffer(),
+                                     (size or 1, maxchar, False))
                     self.assertIs(writer.finish(), ch)
 
     @unittest.skipUnless(support.Py_DEBUG, 'need debug build (Py_DEBUG)')
@@ -2135,60 +2172,83 @@ def test_detect_overflow(self):
     def test_memory_error(self):
         # Inject MemoryError in PyUnicodeWriter_WriteStr()
         writer = self.create_writer(0)
-        writer.write_str("start")
+        writer.write_utf8(b"start", -1)
+        self.assertEqual(writer.get_buffer(), (5, 127, False))
         with self.assertRaises(MemoryError):
             with support.inject_memory_error_cm():
                 # Resize the internal str object
                 writer.write_str("s" * 1024)
         writer.write_str(" end")
         self.assertEqual(writer.finish(), "start end")
 
-        # Inject MemoryError in PyUnicodeWriter_Finish()
+        # Inject MemoryError in PyUnicodeWriter_Finish(). Use write_utf8() to
+        # allocate a buffer of 1024 character. finish() needs to truncate the
+        # buffer to 3 characters.
         writer = self.create_writer(1024)
-        writer.write_str("abc")
+        writer.write_utf8(b"abc", -1)
+        self.assertEqual(writer.get_buffer(), (1024, 127, False))
         with self.assertRaises(MemoryError):
             with support.inject_memory_error_cm():
-                # Need to truncate the internal str object
                 writer.finish()
 
     def test_change_kind(self):
         writer = self.create_writer(0)
+
         # Create an ASCII buffer
         writer.write_str('ascii ')
+        self.assertEqual(writer.get_buffer()[1], 127)
+
         # Change the buffer to UCS1
         writer.write_str('latin1:\xe9 ')
+        self.assertEqual(writer.get_buffer()[1], 255)
+
         # Change the buffer to UCS2
         writer.write_str('ucs2:\u20ac ')
+        self.assertEqual(writer.get_buffer()[1], 0xffff)
+
         # Change the buffer to UCS4
         writer.write_str('ucs4:\U0010ffff')
+        self.assertEqual(writer.get_buffer()[1], 0x10_ffff)
+
         self.assertEqual(writer.finish(),
                          'ascii latin1:\xe9 ucs2:\u20ac ucs4:\U0010ffff')
 
     def test_readonly_optim(self):
         # Read-only optimization: if the first and only write is a Python str
         # object and no buffer was allocated yet, return the object unchanged
         unique_string = 'unique string'
-        writer = self.create_writer(0)
-        writer.write_str(unique_string)
-        self.assertIs(writer.finish(), unique_string)
+        expected = (len(unique_string), 127, True)
+        for size in (0, 123):
+            with self.subTest(size=size):
+                # PyUnicodeWriter_WriteStr() optimization
+                writer = self.create_writer(size)
+                writer.write_str(unique_string)
+                self.assertEqual(writer.get_buffer(), expected)
+                self.assertIs(writer.finish(), unique_string)
 
-        writer = self.create_writer(0)
-        writer.write_substring(unique_string, 0, len(unique_string))
-        self.assertIs(writer.finish(), unique_string)
+                # PyUnicodeWriter_WriteSubstring() optimization
+                writer = self.create_writer(size)
+                writer.write_substring(unique_string, 0, len(unique_string))
+                self.assertEqual(writer.get_buffer(), expected)
+                self.assertIs(writer.finish(), unique_string)
 
-        class MyStr:
-            def __str__(self):
-                return unique_string
-        writer = self.create_writer(0)
-        writer.write_str(MyStr())
-        self.assertIs(writer.finish(), unique_string)
+                # PyUnicodeWriter_WriteStr() optimization
+                class MyStr:
+                    def __str__(self):
+                        return unique_string
+                writer = self.create_writer(size)
+                writer.write_str(MyStr())
+                self.assertEqual(writer.get_buffer(), expected)
+                self.assertIs(writer.finish(), unique_string)
 
-        class MyRepr:
-            def __repr__(self):
-                return unique_string
-        writer = self.create_writer(0)
-        writer.write_repr(MyRepr())
-        self.assertIs(writer.finish(), unique_string)
+                # PyUnicodeWriter_WriteRepr() optimization
+                class MyRepr:
+                    def __repr__(self):
+                        return unique_string
+                writer = self.create_writer(size)
+                writer.write_repr(MyRepr())
+                self.assertEqual(writer.get_buffer(), expected)
+                self.assertIs(writer.finish(), unique_string)
 
 
 # Test PyUnicodeWriter_Format()

Modules/_testcapi/unicode.c

@@ -730,6 +730,31 @@ writer_get_pointer(PyObject *self_raw, PyObject *args)
 }
 
 
+static PyObject*
+writer_get_buffer(PyObject *self_raw, PyObject *args)
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter*)self->writer;
+    PyObject *allocated;
+    Py_UCS4 maxchar;
+    if (writer->buffer) {
+        allocated = PyLong_FromSsize_t(PyUnicode_GET_LENGTH(writer->buffer));
+        maxchar = PyUnicode_MAX_CHAR_VALUE(writer->buffer);
+    }
+    else {
+        allocated = Py_None;
+        maxchar = writer->min_char;
+    }
+    return Py_BuildValue("(NkN)",
+                         allocated, (unsigned long)maxchar,
+                         PyBool_FromLong(writer->readonly));
+}
+
+
 static PyObject*
 writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args))
 {
@@ -755,6 +780,7 @@ static PyMethodDef writer_methods[] = {
     {"write_substring", _PyCFunction_CAST(writer_write_substring), METH_VARARGS},
     {"decodeutf8stateful", _PyCFunction_CAST(writer_decodeutf8stateful), METH_VARARGS},
     {"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},
     {NULL,              NULL}           /* sentinel */
 };

Objects/unicode_writer.c

@@ -151,10 +151,9 @@ PyUnicodeWriter_Create(Py_ssize_t length)
     _PyUnicodeWriter *writer = (_PyUnicodeWriter *)pub_writer;
 
     _PyUnicodeWriter_Init(writer);
-    if (_PyUnicodeWriter_Prepare(writer, length, 127) < 0) {
-        PyUnicodeWriter_Discard(pub_writer);
-        return NULL;
-    }
+    // The buffer is created lazily at the first write, except if
+    // the read-only optimization is used.
+    writer->min_length = length;
     writer->overallocate = 1;
 
     return pub_writer;
@@ -189,9 +188,6 @@ int
 _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
                                  Py_ssize_t length, Py_UCS4 maxchar)
 {
-    Py_ssize_t newlen;
-    PyObject *newbuffer;
-
     assert(length >= 0);
     assert(maxchar <= _Py_MAX_UNICODE);
 
@@ -203,50 +199,50 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
         PyErr_NoMemory();
         return -1;
     }
-    newlen = writer->pos + length;
+    Py_ssize_t alloc = writer->pos + length;
 
     maxchar = Py_MAX(maxchar, writer->min_char);
 
+    PyObject *newbuffer;
     if (writer->buffer == NULL) {
         assert(!writer->readonly);
-        if (writer->overallocate
-            && newlen <= (PY_SSIZE_T_MAX - newlen / OVERALLOCATE_FACTOR)) {
-            /* overallocate to limit the number of realloc() */
-            newlen += newlen / OVERALLOCATE_FACTOR;
-        }
-        if (newlen < writer->min_length)
-            newlen = writer->min_length;
+        // Do not overallocate at the first allocation, but use min_length
+        if (alloc < writer->min_length)
+            alloc = writer->min_length;
 
-        writer->buffer = PyUnicode_New(newlen, maxchar);
+        writer->buffer = PyUnicode_New(alloc, maxchar);
         if (writer->buffer == NULL)
             return -1;
     }
-    else if (newlen > writer->size) {
-        if (writer->overallocate
-            && newlen <= (PY_SSIZE_T_MAX - newlen / OVERALLOCATE_FACTOR)) {
+    else if (alloc > writer->size) {
+        // Do not overallocate at the first allocation, but use min_length
+        int overallocate = (writer->overallocate && !writer->readonly);
+        if (overallocate
+            && alloc <= (PY_SSIZE_T_MAX - alloc / OVERALLOCATE_FACTOR)) {
             /* overallocate to limit the number of realloc() */
-            newlen += newlen / OVERALLOCATE_FACTOR;
+            alloc += alloc / OVERALLOCATE_FACTOR;
         }
-        if (newlen < writer->min_length)
-            newlen = writer->min_length;
+        if (alloc < writer->min_length)
+            alloc = writer->min_length;
 
         if (maxchar > writer->maxchar || writer->readonly) {
             /* resize + widen */
             maxchar = Py_MAX(maxchar, writer->maxchar);
-            newbuffer = PyUnicode_New(newlen, maxchar);
+            newbuffer = PyUnicode_New(alloc, maxchar);
             if (newbuffer == NULL)
                 return -1;
             _PyUnicode_FastCopyCharacters(newbuffer, 0,
                                           writer->buffer, 0, writer->pos);
-            Py_DECREF(writer->buffer);
             writer->readonly = 0;
+            Py_DECREF(writer->buffer);
+            writer->buffer = newbuffer;
         }
         else {
-            newbuffer = _PyUnicode_ResizeCompact(writer->buffer, newlen);
+            newbuffer = _PyUnicode_ResizeCompact(writer->buffer, alloc);
             if (newbuffer == NULL)
                 return -1;
+            writer->buffer = newbuffer;
         }
-        writer->buffer = newbuffer;
     }
     else if (maxchar > writer->maxchar) {
         assert(!writer->readonly);
@@ -310,13 +306,12 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
 {
     assert(PyUnicode_Check(str));
 
-    Py_UCS4 maxchar;
-    Py_ssize_t len;
-
-    len = PyUnicode_GET_LENGTH(str);
-    if (len == 0)
+    Py_ssize_t len = PyUnicode_GET_LENGTH(str);
+    if (len == 0) {
         return 0;
-    maxchar = PyUnicode_MAX_CHAR_VALUE(str);
+    }
+    Py_UCS4 maxchar = PyUnicode_MAX_CHAR_VALUE(str);
+
     if (maxchar > writer->maxchar || len > writer->size - writer->pos) {
         if (writer->buffer == NULL) {
             assert(_PyUnicode_CheckConsistency(str, 1));