diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index f9e7b5ab041962..f13ad6f428ec09 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2050,12 +2050,19 @@ def test_singletons(self): for size in (0, 123): with self.subTest(size=size): writer = self.create_writer(size) + writer.write_utf8(b'utf8', 0) + writer.write_ascii(b'ascii', 0) + writer.write_widechar(b'wstr', 0) + writer.write_ucs4(b'ucs4', 0) + writer.write_substring('text', 0, 0) self.assertIs(writer.finish(), '') for ch in range(256): with self.subTest(ch=ch): ch = chr(ch) writer = self.create_writer(0) + # Use PyUnicodeWriter_WriteSubstring() to avoid the read-only + # buffer optimization writer.write_substring(ch + 'xxx', 0, 1) self.assertIs(writer.finish(), ch) diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c index 7fefd64eefb2ac..91cbfa23e30b20 100644 --- a/Objects/bytesobject.c +++ b/Objects/bytesobject.c @@ -3938,12 +3938,6 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) } } else { - if (size != PyBytes_GET_SIZE(writer->obj)) { - if (_PyBytes_Resize(&writer->obj, size)) { - goto error; - } - } - if (size == 1) { // Get the single byte singleton unsigned char ch = PyBytes_AS_STRING(writer->obj)[0]; @@ -3951,6 +3945,11 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) assert(_Py_IsImmortal(op)); Py_SETREF(writer->obj, op); } + else if (size != PyBytes_GET_SIZE(writer->obj)) { + if (_PyBytes_Resize(&writer->obj, size)) { + goto error; + } + } } result = writer->obj; diff --git a/Objects/unicode_writer.c b/Objects/unicode_writer.c index 92e0db08b9de00..751fca9948598f 100644 --- a/Objects/unicode_writer.c +++ b/Objects/unicode_writer.c @@ -609,8 +609,6 @@ _PyUnicodeWriter_WriteLatin1String(_PyUnicodeWriter *writer, PyObject * _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer) { - PyObject *str; - #ifdef Py_DEBUG // Check for buffer overflow if (writer->buffer != NULL) { @@ -625,23 +623,28 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer) } #endif - if (writer->pos == 0) { - Py_CLEAR(writer->buffer); - return _PyUnicode_GetEmpty(); - } - - str = writer->buffer; + PyObject *str = writer->buffer; writer->buffer = NULL; - if (writer->readonly) { - assert(PyUnicode_GET_LENGTH(str) == writer->pos); - assert(_PyUnicode_CheckConsistency(str, 1)); - return str; + Py_ssize_t final_size = writer->pos; + if (final_size == 0) { + PyObject *empty = _PyUnicode_GetEmpty(); + Py_XDECREF(str); // writer->buffer can be NULL if the position is 0 + return empty; } - if (PyUnicode_GET_LENGTH(str) != writer->pos) { - PyObject *str2; - str2 = _PyUnicode_ResizeCompact(str, writer->pos); + Py_ssize_t length = PyUnicode_GET_LENGTH(str); + if (final_size == 1 && PyUnicode_KIND(str) == PyUnicode_1BYTE_KIND) { + assert(length >= 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; + } + + if (!writer->readonly && length != final_size) { + PyObject *str2 = _PyUnicode_ResizeCompact(str, final_size); if (str2 == NULL) { Py_DECREF(str); return NULL; @@ -650,7 +653,7 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer) } assert(_PyUnicode_CheckConsistency(str, 1)); - return _PyUnicode_Result(str); + return str; }