[pull] main from python:main

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)
 
@@ -2075,6 +2082,39 @@ def test_detect_overflow(self):
         # strategy which depends on the operating system
         self.assertIn(f'at position '.encode(), proc.err)
 
+    @support.nomemtest
+    def test_memory_error(self):
+        # Inject MemoryError in PyUnicodeWriter_WriteStr()
+        writer = self.create_writer(0)
+        writer.write_str("start")
+        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()
+        writer = self.create_writer(1024)
+        writer.write_str("abc")
+        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 ')
+        # Change the buffer to UCS1
+        writer.write_str('latin1:\xe9 ')
+        # Change the buffer to UCS2
+        writer.write_str('ucs2:\u20ac ')
+        # Change the buffer to UCS4
+        writer.write_str('ucs4:\U0010ffff')
+        self.assertEqual(writer.finish(),
+                         'ascii latin1:\xe9 ucs2:\u20ac ucs4:\U0010ffff')
+
 
 # Test PyUnicodeWriter_Format()
 @unittest.skipIf(ctypes is None, 'need ctypes')

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

@@ -199,10 +199,13 @@ void PyUnicodeWriter_Discard(PyUnicodeWriter *writer)
 void
 _PyUnicodeWriter_InitWithBuffer(_PyUnicodeWriter *writer, PyObject *buffer)
 {
+    assert(PyUnstable_Object_IsUniquelyReferenced(buffer));
+
     memset(writer, 0, sizeof(*writer));
     writer->buffer = buffer;
     _PyUnicodeWriter_Update(writer);
     writer->min_length = writer->size;
+    assert(_PyUnicodeWriter_CanWrite(writer));
 }
 
 
@@ -606,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) {
@@ -622,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;
+    }
+
+    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 (PyUnicode_GET_LENGTH(str) != writer->pos) {
-        PyObject *str2;
-        str2 = _PyUnicode_ResizeCompact(str, writer->pos);
+    if (!writer->readonly && length != final_size) {
+        PyObject *str2 = _PyUnicode_ResizeCompact(str, final_size);
         if (str2 == NULL) {
             Py_DECREF(str);
             return NULL;
@@ -647,7 +653,7 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
     }
 
     assert(_PyUnicode_CheckConsistency(str, 1));
-    return _PyUnicode_Result(str);
+    return str;
 }