diff --git a/Doc/c-api/unicode.rst b/Doc/c-api/unicode.rst index 8e93f480d4f7aa4..3efbdfb2c1dc360 100644 --- a/Doc/c-api/unicode.rst +++ b/Doc/c-api/unicode.rst @@ -1813,6 +1813,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 3791b913c175467..b36c65223b27182 100644 --- a/Include/internal/pycore_unicodeobject.h +++ b/Include/internal/pycore_unicodeobject.h @@ -9,8 +9,41 @@ extern "C" { #endif #include "pycore_fileutils.h" // _Py_error_handler +#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_SSIZE_RELAXED #include "pycore_ucnhash.h" // _PyUnicode_Name_CAPI + +// Maximum code point of Unicode 6.0: 0x10ffff (1,114,111). +// The value must be the same in fileutils.c. +#define _Py_MAX_UNICODE 0x10ffff + + +static inline Py_hash_t PyUnicode_HASH(PyObject *op) +{ + assert(PyUnicode_Check(op)); + return FT_ATOMIC_LOAD_SSIZE_RELAXED(_PyASCIIObject_CAST(op)->hash); +} + + +#ifndef NDEBUG +static inline int +_PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer) +{ + // Code adapted from unicode_modifiable() + 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(PyUnicode_HASH(buffer) == -1); + assert(!PyUnicode_CHECK_INTERNED(buffer)); + assert(!_Py_IsImmortal(buffer)); + return 1; +} +#endif + + /* --- Characters Type APIs ----------------------------------------------- */ extern int _PyUnicode_IsXidStart(Py_UCS4 ch); diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index dc6204d0a84658b..b890f71ea29781d 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -3361,3 +3361,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 95ef7e7461a5d46..20deb0446cb581a 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2088,10 +2088,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): @@ -2100,6 +2098,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 5182615ff2d7949..5c830cd330acc25 100644 --- a/Objects/longobject.c +++ b/Objects/longobject.c @@ -2144,6 +2144,7 @@ long_to_decimal_string_internal(PyObject *aa, Py_DECREF(scratch); return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); } else if (bytes_writer) { *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, strlen); @@ -2313,8 +2314,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_Prepare(bytes_writer, *bytes_str, sz); diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c index 64c0a912f5ade9f..56ddb572e6b781d 100644 --- a/Objects/unicodeobject.c +++ b/Objects/unicodeobject.c @@ -102,9 +102,7 @@ NOTE: In the interpreter's initialization phase, some globals are currently */ -// Maximum code point of Unicode 6.0: 0x10ffff (1,114,111). -// The value must be the same in fileutils.c. -#define MAX_UNICODE 0x10ffff +#define MAX_UNICODE _Py_MAX_UNICODE #ifdef Py_DEBUG # define _PyUnicode_CHECK(op) _PyUnicode_CheckConsistency(op, 0) @@ -166,12 +164,6 @@ static inline void PyUnicode_SET_UTF8_LENGTH(PyObject *op, Py_ssize_t length) #define _PyUnicode_HASH(op) \ (_PyASCIIObject_CAST(op)->hash) -static inline Py_hash_t PyUnicode_HASH(PyObject *op) -{ - assert(_PyUnicode_CHECK(op)); - return FT_ATOMIC_LOAD_SSIZE_RELAXED(_PyASCIIObject_CAST(op)->hash); -} - static inline void PyUnicode_SET_HASH(PyObject *op, Py_hash_t hash) { FT_ATOMIC_STORE_SSIZE_RELAXED(_PyASCIIObject_CAST(op)->hash, hash); @@ -1841,18 +1833,21 @@ unicode_is_singleton(PyObject *unicode) } #endif +// If this function is updated, update also _PyUnicodeWriter_CanWrite(). static int unicode_modifiable(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)); @@ -2156,6 +2151,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; @@ -2453,6 +2449,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; @@ -2739,8 +2736,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)) { @@ -3030,8 +3029,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) @@ -5587,6 +5588,7 @@ unicode_decode_utf8_writer(_PyUnicodeWriter *writer, if (_PyUnicodeWriter_Prepare(writer, size, 127) < 0) { return -1; } + assert(_PyUnicodeWriter_CanWrite(writer)); const char *starts = s; const char *end = s + size; @@ -13854,10 +13856,13 @@ void PyUnicodeWriter_Discard(PyUnicodeWriter *writer) static inline 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)); } @@ -13966,6 +13971,7 @@ _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch) assert(ch <= 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; @@ -14013,6 +14019,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; @@ -14090,6 +14098,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); @@ -14147,8 +14156,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) { @@ -14252,6 +14263,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;