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..fcdf75ad50d74c 100644 --- a/Include/internal/pycore_unicodeobject.h +++ b/Include/internal/pycore_unicodeobject.h @@ -104,12 +104,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/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 06c13dcfe75720..30549863cd04fd 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): @@ -2113,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/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..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)); } @@ -350,6 +353,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 +433,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 +491,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 +599,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;