From e42d8743a04596985c178336169bcd069a3aa83f Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Fri, 18 Sep 2026 05:07:47 +0200
Subject: [PATCH] gh-157710: Add MemoryError tests to PyUnicodeWriter

Add a test changing the buffer kind multiple times.

Add assertions to _PyUnicodeWriter_InitWithBuffer().
---
 Lib/test/test_capi/test_unicode.py | 33 ++++++++++++++++++++++++++++++
 Objects/unicode_writer.c           |  3 +++
 2 files changed, 36 insertions(+)

diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py
index 50e807e3e7db08..f9e7b5ab041962 100644
--- a/Lib/test/test_capi/test_unicode.py
+++ b/Lib/test/test_capi/test_unicode.py
@@ -2075,6 +2075,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/unicode_writer.c b/Objects/unicode_writer.c
index fe1bd97775b3ae..92e0db08b9de00 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));
 }