[3.14] gh-157710: Add MemoryError tests to PyUnicodeWriter

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.

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);

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()

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')

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);

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;