gh-158439: Add _PyUnicodeWriter.recheck_maxchar (#158442)

If PyUnicodeWriter_WriteUTF8() or PyUnicodeWriter_Format() fails,
PyUnicodeWriter_Finish() now checks if the string uses the most
efficient storage. Previously, PyUnicodeWriter_Finish() could return
a string with an inefficient storage.

* Add _PyUnicodeWriter.recheck_maxchar. PyUnicodeWriter_WriteUTF8()
  and PyUnicodeWriter_Format() now set recheck_maxchar to 1 on error.
* Move error handling from caller to _PyUnicode_DecodeUTF8Writer().
* PyUnicodeWriter_DecodeUTF8Stateful(): rename length parameter to
  size and rename string parameter to str.

Author: vstinner

Doc/c-api/unicode.rst

@@ -1945,7 +1945,7 @@ object.
    On success, return ``0``.
    On error, set an exception, leave the writer unchanged, and return ``-1``.
 
-.. c:function:: int PyUnicodeWriter_DecodeUTF8Stateful(PyUnicodeWriter *writer, const char *string, Py_ssize_t length, const char *errors, Py_ssize_t *consumed)
+.. c:function:: int PyUnicodeWriter_DecodeUTF8Stateful(PyUnicodeWriter *writer, const char *str, Py_ssize_t size, const char *errors, Py_ssize_t *consumed)
 
    Decode the string *str* from UTF-8 with *errors* error handler and write the
    output into *writer*.

Include/cpython/unicodeobject.h

@@ -516,8 +516,8 @@ PyAPI_FUNC(int) PyUnicodeWriter_Format(
     ...);
 PyAPI_FUNC(int) PyUnicodeWriter_DecodeUTF8Stateful(
     PyUnicodeWriter *writer,
-    const char *string,         /* UTF-8 encoded string */
-    Py_ssize_t length,          /* size of string */
+    const char *str,            /* UTF-8 encoded string */
+    Py_ssize_t size,            /* size of string */
     const char *errors,         /* error handling */
     Py_ssize_t *consumed);      /* bytes consumed */
 
@@ -538,6 +538,9 @@ typedef struct {
     /* minimum character (default: 127, ASCII) */
     Py_UCS4 min_char;
 
+    // If non-zero, _PyUnicodeWriter_Finish() needs to check maxchar.
+    int recheck_maxchar;
+
     /* If non-zero, overallocate the buffer (default: 0). */
     unsigned char overallocate;
 

Lib/test/test_capi/test_unicode.py

@@ -1976,7 +1976,11 @@ def test_invalid_utf8(self):
             writer.write_utf8(b"invalid=\xFF", -1)
 
     def test_recover_utf8_error(self):
-        # test recovering from PyUnicodeWriter_WriteUTF8() error
+        # Recover from PyUnicodeWriter_WriteUTF8() errors. A temporary write
+        # changes the buffer kind to UCS-2 before raising UnicodeDecodeError.
+        # Then, PyUnicodeWriter_Finish() has to change the buffer kind back to
+        # ASCII.
+
         writer = self.create_writer(0)
         writer.write_utf8(b"value=", -1)
 
@@ -1986,12 +1990,13 @@ def test_recover_utf8_error(self):
         with self.assertRaises(UnicodeDecodeError):
             s = "truncated\u20AC".encode()
             writer.write_utf8(s, len(s) - 1)
+        with self.assertRaises(UnicodeDecodeError):
+            # Change buffer kind to UCS-2 then raise UnicodeDecodeError
+            s = "\u20AC\u20AC".encode()
+            writer.write_utf8(s, len(s) - 1)
 
-        # retry write with a valid string
         writer.write_utf8(b"valid", -1)
-
-        self.assertEqual(writer.finish(),
-                         "value=valid")
+        self.assertEqual(writer.finish(), "value=valid")
 
     def test_decode_utf8(self):
         # test PyUnicodeWriter_DecodeUTF8Stateful()
@@ -2299,6 +2304,29 @@ def test_recover_error(self):
 
         self.assertEqual(writer.finish(), 'Hello World.')
 
+    def test_recheck_maxchar(self):
+        # PyUnicodeWriter_Format() changes buffer kind to UCS-2 before raising
+        # an exception. Then, PyUnicodeWriter_Finish() has to change the buffer
+        # kind back to ASCII.
+        from ctypes import py_object
+
+        class StrError:
+            def __str__(self):
+                raise RuntimeError("bug")
+
+        writer = self.create_writer(0)
+        # Allocate ASCII buffer
+        writer.write_str('ascii')
+
+        obj = StrError()
+        ucs2_utf8 = '\u20ac'.encode()
+        with self.assertRaises(RuntimeError):
+            # Change buffer kind to UCS-2, but then raise RuntimeError
+            self.writer_format(writer, b"%s%S", ucs2_utf8, py_object(obj))
+
+        writer.write_str('.')
+        self.assertEqual(writer.finish(), 'ascii.')
+
     def test_readonly_optim(self):
         # Read-only optimization: if the first and only write is a Python str
         # object and no buffer was allocated yet, return the object unchanged

Misc/NEWS.d/next/C_API/2026-09-29-19-01-56.gh-issue-158439.ybKo_6.rst

@@ -0,0 +1,4 @@
+If :c:func:`PyUnicodeWriter_WriteUTF8` or :c:func:`PyUnicodeWriter_Format`
+fails, :c:func:`PyUnicodeWriter_Finish` now checks if the string uses the most
+efficient storage. Previously, :c:func:`PyUnicodeWriter_Finish` could return a
+string with an inefficient storage. Patch by Victor Stinner.

Objects/unicode_writer.c

@@ -160,7 +160,8 @@ PyUnicodeWriter_Create(Py_ssize_t length)
 }
 
 
-void PyUnicodeWriter_Discard(PyUnicodeWriter *writer)
+void
+PyUnicodeWriter_Discard(PyUnicodeWriter *writer)
 {
     if (writer == NULL) {
         return;
@@ -511,8 +512,7 @@ PyUnicodeWriter_WriteASCII(PyUnicodeWriter *writer,
     assert(writer != NULL);
     _Py_AssertHoldsTstate();
 
-    _PyUnicodeWriter *priv_writer = (_PyUnicodeWriter*)writer;
-    return _PyUnicodeWriter_WriteASCIIString(priv_writer, str, size);
+    return _PyUnicodeWriter_WriteASCIIString((_PyUnicodeWriter*)writer, str, size);
 }
 
 
@@ -525,40 +525,25 @@ PyUnicodeWriter_WriteUTF8(PyUnicodeWriter *writer,
         size = strlen(str);
     }
 
-    _PyUnicodeWriter *_writer = (_PyUnicodeWriter*)writer;
-    Py_ssize_t old_pos = _writer->pos;
-    int res = _PyUnicode_DecodeUTF8Writer(_writer, str, size,
-                                          _Py_ERROR_STRICT, NULL, NULL);
-    if (res < 0) {
-        _writer->pos = old_pos;
-    }
-    return res;
+    return _PyUnicode_DecodeUTF8Writer((_PyUnicodeWriter*)writer, str, size,
+                                       _Py_ERROR_STRICT, NULL, NULL);
 }
 
 
 int
 PyUnicodeWriter_DecodeUTF8Stateful(PyUnicodeWriter *writer,
-                                   const char *string,
-                                   Py_ssize_t length,
+                                   const char *str,
+                                   Py_ssize_t size,
                                    const char *errors,
                                    Py_ssize_t *consumed)
 {
-    if (length < 0) {
-        length = strlen(string);
-    }
-
-    _PyUnicodeWriter *_writer = (_PyUnicodeWriter*)writer;
-    Py_ssize_t old_pos = _writer->pos;
-    int res = _PyUnicode_DecodeUTF8Writer(_writer, string, length,
-                                          _Py_ERROR_UNKNOWN, errors,
-                                          consumed);
-    if (res < 0) {
-        _writer->pos = old_pos;
-        if (consumed) {
-            *consumed = 0;
-        }
+    if (size < 0) {
+        size = strlen(str);
     }
-    return res;
+
+    return _PyUnicode_DecodeUTF8Writer((_PyUnicodeWriter*)writer, str, size,
+                                       _Py_ERROR_UNKNOWN, errors,
+                                       consumed);
 }
 
 
@@ -600,30 +585,55 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
 
     Py_ssize_t final_size = writer->pos;
     if (final_size == 0) {
+        // Get the empty string singleton
         PyObject *empty = _PyUnicode_GetEmpty();
         Py_XDECREF(str);  // writer->buffer can be NULL if the position is 0
         return empty;
     }
 
-    Py_ssize_t length = PyUnicode_GET_LENGTH(str);
+    if (writer->readonly) {
+        assert(final_size == PyUnicode_GET_LENGTH(str));
+        goto done;
+    }
+    assert(final_size <= PyUnicode_GET_LENGTH(str));
+
     if (final_size == 1 && PyUnicode_KIND(str) == PyUnicode_1BYTE_KIND) {
-        assert(length >= 1);
+        // Get the single character singleton
+        assert(PyUnicode_GET_LENGTH(str) >= 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;
+        str = _Py_LATIN1_CHR(ch);
+        goto done;
+    }
+
+    if (writer->recheck_maxchar) {
+        Py_UCS4 maxchar = _PyUnicode_FindMaxChar(str, 0, final_size);
+        if (maxchar != writer->maxchar) {
+            // Adjust the string kind
+            PyObject *str2 = PyUnicode_New(final_size, maxchar);
+            if (str2 == NULL) {
+                Py_DECREF(str);
+                return NULL;
+            }
+            _PyUnicode_FastCopyCharacters(str2, 0, str, 0, final_size);
+            Py_SETREF(str, str2);
+            goto done;
+        }
     }
 
-    if (!writer->readonly && length != final_size) {
+    if (PyUnicode_GET_LENGTH(str) != final_size) {
+        // Truncate the string
         PyObject *str2 = _PyUnicode_ResizeCompact(str, final_size);
         if (str2 == NULL) {
             Py_DECREF(str);
             return NULL;
         }
         str = str2;
+        goto done;
     }
 
+done:
     assert(_PyUnicode_CheckConsistency(str, 1));
     return str;
 }

Objects/unicodeobject.c

@@ -3166,16 +3166,21 @@ PyUnicodeWriter_Format(PyUnicodeWriter *writer, const char *format, ...)
 }
 
 int
-_PyUnicodeWriter_FormatV(PyUnicodeWriter *writer, const char *format,
+_PyUnicodeWriter_FormatV(PyUnicodeWriter *pub_writer, const char *format,
                          va_list vargs)
 {
-    _PyUnicodeWriter *_writer = (_PyUnicodeWriter*)writer;
-    Py_ssize_t old_pos = _writer->pos;
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter*)pub_writer;
+    Py_ssize_t old_pos = writer->pos;
+    Py_UCS4 old_maxchar = writer->maxchar;
 
-    int res = unicode_from_format(_writer, format, vargs);
+    int res = unicode_from_format(writer, format, vargs);
 
     if (res < 0) {
-        _writer->pos = old_pos;
+        writer->pos = old_pos;
+        if (writer->maxchar > old_maxchar) {
+            // _PyUnicodeWriter_Finish() will check maxchar
+            writer->recheck_maxchar = 1;
+        }
     }
     return res;
 }
@@ -5364,7 +5369,7 @@ unicode_decode_utf8(const char *s, Py_ssize_t size,
 }
 
 
-// Used by PyUnicodeWriter_WriteUTF8() implementation
+// Used by PyUnicodeWriter_WriteUTF8() and PyUnicodeWriter_DecodeUTF8Stateful()
 int
 _PyUnicode_DecodeUTF8Writer(_PyUnicodeWriter *writer,
                             const char *s, Py_ssize_t size,
@@ -5378,18 +5383,20 @@ _PyUnicode_DecodeUTF8Writer(_PyUnicodeWriter *writer,
         return 0;
     }
 
+    Py_ssize_t old_pos = writer->pos;
+    Py_UCS4 old_maxchar = writer->maxchar;
+
     // fast path: try ASCII string.
     if (_PyUnicodeWriter_Prepare(writer, size, 127) < 0) {
-        return -1;
+        goto error;
     }
     assert(_PyUnicodeWriter_CanWrite(writer));
 
     const char *starts = s;
     const char *end = s + size;
-    Py_ssize_t decoded = 0;
-    Py_UCS1 *dest = (Py_UCS1*)writer->data + writer->pos * writer->kind;
     if (writer->kind == PyUnicode_1BYTE_KIND) {
-        decoded = ascii_decode(s, end, dest);
+        Py_UCS1 *dest = (Py_UCS1*)writer->data + writer->pos * writer->kind;
+        Py_ssize_t decoded = ascii_decode(s, end, dest);
         writer->pos += decoded;
 
         if (decoded == size) {
@@ -5401,8 +5408,24 @@ _PyUnicode_DecodeUTF8Writer(_PyUnicodeWriter *writer,
         s += decoded;
     }
 
-    return unicode_decode_utf8_impl(writer, starts, s, end,
-                                    error_handler, errors, consumed);
+    int res = unicode_decode_utf8_impl(writer, starts, s, end,
+                                       error_handler, errors, consumed);
+    if (res < 0) {
+        goto error;
+    }
+    return 0;
+
+error:
+    // Restore the writer to its previous state
+    writer->pos = old_pos;
+    if (writer->maxchar > old_maxchar) {
+        // _PyUnicodeWriter_Finish() will check maxchar
+        writer->recheck_maxchar = 1;
+    }
+    if (consumed) {
+        *consumed = 0;
+    }
+    return -1;
 }