diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index 50e807e3e7db083..f13ad6f428ec095 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) @@ -2075,6 +2082,39 @@ def test_detect_overflow(self): # strategy which depends on the operating system self.assertIn(f'at position '.encode(), proc.err) + @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/bytesobject.c b/Objects/bytesobject.c index 7fefd64eefb2acd..91cbfa23e30b208 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 fe1bd97775b3ae2..751fca9948598ff 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)); } @@ -606,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) { @@ -622,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; + } + + 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 (PyUnicode_GET_LENGTH(str) != writer->pos) { - PyObject *str2; - str2 = _PyUnicode_ResizeCompact(str, writer->pos); + if (!writer->readonly && length != final_size) { + PyObject *str2 = _PyUnicode_ResizeCompact(str, final_size); if (str2 == NULL) { Py_DECREF(str); return NULL; @@ -647,7 +653,7 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer) } assert(_PyUnicode_CheckConsistency(str, 1)); - return _PyUnicode_Result(str); + return str; }