From 6b974528366e5452300c1f3de27f5062b1160039 Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Sun, 20 Sep 2026 14:01:40 +0200 Subject: [PATCH 1/2] gh-157710: Add MemoryError tests to PyUnicodeWriter (#157853) Add a test changing the buffer kind multiple times. Add assertions to _PyUnicodeWriter_InitWithBuffer(). --- Lib/test/test_capi/test_unicode.py | 33 ++++++++++++++++++++++++++++++ Objects/unicode_writer.c | 3 +++ 2 files changed, 36 insertions(+) diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index 50e807e3e7db083..f9e7b5ab0419621 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2075,6 +2075,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') diff --git a/Objects/unicode_writer.c b/Objects/unicode_writer.c index fe1bd97775b3ae2..92e0db08b9de002 100644 --- a/Objects/unicode_writer.c +++ b/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)); } From a64b7ecba4db5480504a1a83ae6e80c81472fe16 Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Sun, 20 Sep 2026 16:25:20 +0200 Subject: [PATCH 2/2] 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. --- Lib/test/test_capi/test_unicode.py | 7 ++++++ Objects/bytesobject.c | 11 +++++----- Objects/unicode_writer.c | 35 ++++++++++++++++-------------- 3 files changed, 31 insertions(+), 22 deletions(-) diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index f9e7b5ab0419621..f13ad6f428ec095 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/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) diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c index 7fefd64eefb2acd..91cbfa23e30b208 100644 --- a/Objects/bytesobject.c +++ b/Objects/bytesobject.c @@ -3938,12 +3938,6 @@ 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]; @@ -3951,6 +3945,11 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) 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; diff --git a/Objects/unicode_writer.c b/Objects/unicode_writer.c index 92e0db08b9de002..751fca9948598ff 100644 --- a/Objects/unicode_writer.c +++ b/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; }