gh-157710: Avoid resize in PyUnicodeWriter_Finish() for singleton (#157862)

PyBytesWriter_FinishWithSize() and PyUnicodeWriter_Finish() now
discard the output string instead of resizing it if the single
byte/character singleton is used.

Author: vstinner

Lib/test/test_capi/test_unicode.py

@@ -2050,12 +2050,19 @@ def test_singletons(self):
         for size in (0, 123):
             with self.subTest(size=size):
                 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)
                 self.assertIs(writer.finish(), '')
 
         for ch in range(256):
             with self.subTest(ch=ch):
                 ch = chr(ch)
                 writer = self.create_writer(0)
+                # Use PyUnicodeWriter_WriteSubstring() to avoid the read-only
+                # buffer optimization
                 writer.write_substring(ch + 'xxx', 0, 1)
                 self.assertIs(writer.finish(), ch)
 

Objects/bytesobject.c

@@ -3938,19 +3938,18 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
             }
         }
         else {
-            if (size != PyBytes_GET_SIZE(writer->obj)) {
-                if (_PyBytes_Resize(&writer->obj, size)) {
-                    goto error;
-                }
-            }
-
             if (size == 1) {
                 // Get the single byte singleton
                 unsigned char ch = PyBytes_AS_STRING(writer->obj)[0];
                 PyObject *op = (PyObject*)CHARACTER(ch);
                 assert(_Py_IsImmortal(op));
                 Py_SETREF(writer->obj, op);
             }
+            else if (size != PyBytes_GET_SIZE(writer->obj)) {
+                if (_PyBytes_Resize(&writer->obj, size)) {
+                    goto error;
+                }
+            }
         }
 
         result = writer->obj;

Objects/unicode_writer.c

@@ -609,8 +609,6 @@ _PyUnicodeWriter_WriteLatin1String(_PyUnicodeWriter *writer,
 PyObject *
 _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
 {
-    PyObject *str;
-
 #ifdef Py_DEBUG
     // Check for buffer overflow
     if (writer->buffer != NULL) {
@@ -625,23 +623,28 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
     }
 #endif
 
-    if (writer->pos == 0) {
-        Py_CLEAR(writer->buffer);
-        return _PyUnicode_GetEmpty();
-    }
-
-    str = writer->buffer;
+    PyObject *str = writer->buffer;
     writer->buffer = NULL;
 
-    if (writer->readonly) {
-        assert(PyUnicode_GET_LENGTH(str) == writer->pos);
-        assert(_PyUnicode_CheckConsistency(str, 1));
-        return str;
+    Py_ssize_t final_size = writer->pos;
+    if (final_size == 0) {
+        PyObject *empty = _PyUnicode_GetEmpty();
+        Py_XDECREF(str);  // writer->buffer can be NULL if the position is 0
+        return empty;
     }
 
-    if (PyUnicode_GET_LENGTH(str) != writer->pos) {
-        PyObject *str2;
-        str2 = _PyUnicode_ResizeCompact(str, writer->pos);
+    Py_ssize_t length = PyUnicode_GET_LENGTH(str);
+    if (final_size == 1 && PyUnicode_KIND(str) == PyUnicode_1BYTE_KIND) {
+        assert(length >= 1);
+        const Py_UCS1 *data = PyUnicode_1BYTE_DATA(str);
+        Py_UCS1 ch = data[0];
+        PyObject *latin1_char = _Py_LATIN1_CHR(ch);
+        Py_DECREF(str);
+        return latin1_char;
+    }
+
+    if (!writer->readonly && length != final_size) {
+        PyObject *str2 = _PyUnicode_ResizeCompact(str, final_size);
         if (str2 == NULL) {
             Py_DECREF(str);
             return NULL;
@@ -650,7 +653,7 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
     }
 
     assert(_PyUnicode_CheckConsistency(str, 1));
-    return _PyUnicode_Result(str);
+    return str;
 }