From 6d9d6aa358378b4d0f12c9dcfd6eec2ffe691a2d Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Fri, 18 Sep 2026 02:32:50 +0200 Subject: [PATCH 1/3] gh-157710: Add _PyUnicodeWriter_CanWrite() function (#157712) Make sure that a writer can be modified before writing into it. * Document that PyUnicodeWriter is not thread safe. * Add singletons tests to test_capi.test_unicode. * Check PyUnicode_CheckExact() earlier in _PyUnicode_IsModifiable(). (cherry picked from commit 854295809ad5a42c9461f34d47b1252a3b4a2027) --- Doc/c-api/unicode.rst | 3 +++ Include/internal/pycore_unicodeobject.h | 23 ++++++++++++++++++++++- Lib/test/test_capi/test_unicode.py | 6 ++---- Objects/longobject.c | 5 ++++- Objects/unicode_writer.c | 8 +++++++- Objects/unicodeobject.c | 18 ++++++++++++++---- 6 files changed, 52 insertions(+), 11 deletions(-) diff --git a/Doc/c-api/unicode.rst b/Doc/c-api/unicode.rst index 268bcffc5dd1d1..5b3a1b958a7217 100644 --- a/Doc/c-api/unicode.rst +++ b/Doc/c-api/unicode.rst @@ -1826,6 +1826,9 @@ object. The instance must be destroyed by :c:func:`PyUnicodeWriter_Finish` on success, or :c:func:`PyUnicodeWriter_Discard` on error. + The API is **not thread safe**. To share a writer with multiple threads, a + critical section or a lock is needed. + .. c:function:: PyUnicodeWriter* PyUnicodeWriter_Create(Py_ssize_t length) Create a Unicode writer instance. diff --git a/Include/internal/pycore_unicodeobject.h b/Include/internal/pycore_unicodeobject.h index 74d84052a2bb2b..62ef99aeb39806 100644 --- a/Include/internal/pycore_unicodeobject.h +++ b/Include/internal/pycore_unicodeobject.h @@ -16,7 +16,9 @@ extern "C" { #define _Py_MAX_UNICODE 0x10ffff -extern int _PyUnicode_IsModifiable(PyObject *unicode); +// Export for '_multibytecodec' shared extension. _PyUnicodeWriter_CanWrite() +// calls this function when assertions are enabled. +PyAPI_FUNC(int) _PyUnicode_IsModifiable(PyObject *unicode); extern void _PyUnicodeWriter_InitWithBuffer( _PyUnicodeWriter *writer, PyObject *buffer); @@ -104,12 +106,31 @@ _PyUnicode_EnsureUnicode(PyObject *obj) return 0; } +#ifndef NDEBUG +static inline int +_PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer) +{ + // Code adapted from _PyUnicode_IsModifiable() + assert(!writer->readonly); + PyObject *buffer = writer->buffer; + assert(buffer != NULL); + // Do not use _PyObject_IsUniquelyReferenced(): the caller can have its own + // lock to prevent a writer being used by two theads at the same time. + assert(Py_REFCNT(buffer) == 1); + assert(PyUnstable_Unicode_GET_CACHED_HASH(buffer) == -1); + assert(!PyUnicode_CHECK_INTERNED(buffer)); + assert(!_Py_IsImmortal(buffer)); + return 1; +} +#endif + static inline int _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch) { assert(ch <= _Py_MAX_UNICODE); if (_PyUnicodeWriter_Prepare(writer, 1, ch) < 0) return -1; + assert(_PyUnicodeWriter_CanWrite(writer)); PyUnicode_WRITE(writer->kind, writer->data, writer->pos, ch); writer->pos++; return 0; diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index 06c13dcfe75720..c45628c12a3268 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2101,10 +2101,8 @@ def test_substring_empty(self): self.assertEqual(writer.finish(), '') def test_singletons(self): - for size in (0, 123): - with self.subTest(size=size): - writer = self.create_writer(size) - self.assertIs(writer.finish(), '') + writer = self.create_writer(5) + self.assertIs(writer.finish(), '') for ch in range(256): with self.subTest(ch=ch): diff --git a/Objects/longobject.c b/Objects/longobject.c index 6454565aebf6a1..ec4368b87a247d 100644 --- a/Objects/longobject.c +++ b/Objects/longobject.c @@ -2220,6 +2220,7 @@ long_to_decimal_string_internal(PyObject *aa, Py_DECREF(scratch); return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); } else if (bytes_writer) { *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, strlen, @@ -2390,8 +2391,10 @@ long_format_binary(PyObject *aa, int base, int alternate, } if (writer) { - if (_PyUnicodeWriter_Prepare(writer, sz, 'x') == -1) + if (_PyUnicodeWriter_Prepare(writer, sz, 'x') == -1) { return -1; + } + assert(_PyUnicodeWriter_CanWrite(writer)); } else if (bytes_writer) { *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, sz, diff --git a/Objects/unicode_writer.c b/Objects/unicode_writer.c index a753c9b971c702..b10d9e94098935 100644 --- a/Objects/unicode_writer.c +++ b/Objects/unicode_writer.c @@ -350,6 +350,8 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str) if (_PyUnicodeWriter_PrepareInternal(writer, len, maxchar) == -1) return -1; } + + assert(_PyUnicodeWriter_CanWrite(writer)); _PyUnicode_FastCopyCharacters(writer->buffer, writer->pos, str, 0, len); writer->pos += len; @@ -428,6 +430,7 @@ _PyUnicodeWriter_WriteSubstring(_PyUnicodeWriter *writer, PyObject *str, if (_PyUnicodeWriter_Prepare(writer, len, maxchar) < 0) { return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); _PyUnicode_FastCopyCharacters(writer->buffer, writer->pos, str, start, len); @@ -485,8 +488,10 @@ _PyUnicodeWriter_WriteASCIIString(_PyUnicodeWriter *writer, return 0; } - if (_PyUnicodeWriter_Prepare(writer, len, 127) == -1) + if (_PyUnicodeWriter_Prepare(writer, len, 127) == -1) { return -1; + } + assert(_PyUnicodeWriter_CanWrite(writer)); switch (writer->kind) { @@ -591,6 +596,7 @@ _PyUnicodeWriter_WriteLatin1String(_PyUnicodeWriter *writer, maxchar = ucs1lib_find_max_char((const Py_UCS1*)str, (const Py_UCS1*)str + len); if (_PyUnicodeWriter_Prepare(writer, len, maxchar) == -1) return -1; + assert(_PyUnicodeWriter_CanWrite(writer)); unicode_write_cstr(writer->buffer, writer->pos, str, len); writer->pos += len; return 0; diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c index 8fbe8a714ff7c6..df7e4e28bbc25a 100644 --- a/Objects/unicodeobject.c +++ b/Objects/unicodeobject.c @@ -1743,18 +1743,21 @@ unicode_is_singleton(PyObject *unicode) } #endif +// If this function is updated, update also _PyUnicodeWriter_CanWrite(). int _PyUnicode_IsModifiable(PyObject *unicode) { assert(_PyUnicode_CHECK(unicode)); + if (!PyUnicode_CheckExact(unicode)) + return 0; + // On Free Threading, this test fails if called from a thread other + // than the one which created the str object. if (!_PyObject_IsUniquelyReferenced(unicode)) return 0; if (PyUnicode_HASH(unicode) != -1) return 0; if (PyUnicode_CHECK_INTERNED(unicode)) return 0; - if (!PyUnicode_CheckExact(unicode)) - return 0; #ifdef Py_DEBUG /* singleton refcount is greater than 1 */ assert(!unicode_is_singleton(unicode)); @@ -2006,6 +2009,7 @@ PyUnicodeWriter_WriteWideChar(PyUnicodeWriter *pub_writer, if (_PyUnicodeWriter_Prepare(writer, size - num_surrogates, maxchar) < 0) { return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); int kind = writer->kind; void *data = (Py_UCS1*)writer->data + writer->pos * kind; @@ -2303,6 +2307,7 @@ PyUnicodeWriter_WriteUCS4(PyUnicodeWriter *pub_writer, if (_PyUnicodeWriter_Prepare(writer, size, maxchar) < 0) { return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); int kind = writer->kind; void *data = (Py_UCS1*)writer->data + writer->pos * kind; @@ -2589,8 +2594,10 @@ unicode_fromformat_write_str(_PyUnicodeWriter *writer, PyObject *str, else maxchar = writer->maxchar; - if (_PyUnicodeWriter_Prepare(writer, arglen, maxchar) == -1) + if (_PyUnicodeWriter_Prepare(writer, arglen, maxchar) == -1) { return -1; + } + assert(_PyUnicodeWriter_CanWrite(writer)); fill = Py_MAX(width - length, 0); if (fill && !(flags & F_LJUST)) { @@ -2880,8 +2887,10 @@ unicode_fromformat_arg(_PyUnicodeWriter *writer, Py_ssize_t spacepad = Py_MAX(width - precision - sign, 0); Py_ssize_t zeropad = Py_MAX(precision - len, 0); - if (_PyUnicodeWriter_Prepare(writer, width, 127) == -1) + if (_PyUnicodeWriter_Prepare(writer, width, 127) == -1) { return NULL; + } + assert(_PyUnicodeWriter_CanWrite(writer)); if (spacepad && !(flags & F_LJUST)) { if (PyUnicode_Fill(writer->buffer, writer->pos, spacepad, ' ') == -1) @@ -5408,6 +5417,7 @@ _PyUnicode_DecodeUTF8Writer(_PyUnicodeWriter *writer, if (_PyUnicodeWriter_Prepare(writer, size, 127) < 0) { return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); const char *starts = s; const char *end = s + size; From ecffa9b4b6cdeccb8cad2ddee1b9da52f73c1de9 Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Sun, 20 Sep 2026 14:01:40 +0200 Subject: [PATCH 2/3] gh-157710: Add MemoryError tests to PyUnicodeWriter (#157853) Add a test changing the buffer kind multiple times. Add assertions to _PyUnicodeWriter_InitWithBuffer(). (cherry picked from commit 6b974528366e5452300c1f3de27f5062b1160039) --- Lib/test/support/__init__.py | 19 +++++++++++++++++ Lib/test/test_capi/test_unicode.py | 33 ++++++++++++++++++++++++++++++ Objects/unicode_writer.c | 3 +++ 3 files changed, 55 insertions(+) diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index cbfbac23aafdbb..4d09bd7573d0e6 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -3453,3 +3453,22 @@ def built_with_c_assertions(): return Py_DEBUG else: return bool(_testlimitedcapi._py_getbuiltwithassert()) + + +@contextlib.contextmanager +def inject_memory_error_cm(start=0, stop=0): + """ + Similar to inject_memory_error() but can be used as a context manager. + + Raise SkipTest if the _testcapi extension module is missing + """ + try: + import _testcapi + except ImportError: + raise unittest.SkipTest("_testcapi required") + + try: + _testcapi.set_nomemory(start, stop) + yield + finally: + _testcapi.remove_mem_hooks() diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index c45628c12a3268..30549863cd04fd 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2111,6 +2111,39 @@ def test_singletons(self): writer.write_substring(ch + 'xxx', 0, 1) self.assertIs(writer.finish(), ch) + @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 b10d9e94098935..58b5d60333a786 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 c552462ea3d6d9755398a2076fa8d1e2a7c242ed Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Fri, 16 Oct 2026 17:45:35 +0200 Subject: [PATCH 3/3] No need to export _PyUnicode_IsModifiable() --- Include/internal/pycore_unicodeobject.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/Include/internal/pycore_unicodeobject.h b/Include/internal/pycore_unicodeobject.h index 62ef99aeb39806..fcdf75ad50d74c 100644 --- a/Include/internal/pycore_unicodeobject.h +++ b/Include/internal/pycore_unicodeobject.h @@ -16,9 +16,7 @@ extern "C" { #define _Py_MAX_UNICODE 0x10ffff -// Export for '_multibytecodec' shared extension. _PyUnicodeWriter_CanWrite() -// calls this function when assertions are enabled. -PyAPI_FUNC(int) _PyUnicode_IsModifiable(PyObject *unicode); +extern int _PyUnicode_IsModifiable(PyObject *unicode); extern void _PyUnicodeWriter_InitWithBuffer( _PyUnicodeWriter *writer, PyObject *buffer);