gh-160324: Fix io.StringIO state after MemoryError

Lib/test/test_io/test_memoryio.py

@@ -1158,6 +1158,39 @@ def __str__(self):
         memio2.write(MyStr("world"))
         self.assertEqual(memio2.getvalue(), "hello world")
 
+    @support.nomemtest
+    def test_memory_error_consistent_state(self):
+        # A MemoryError while copying the accumulated writes into the
+        # internal buffer must leave the size of the stream consistent with
+        # the buffer, so that using the object again does not read or write
+        # outside it.
+        piece = "0123456789" * 10
+        text = piece * 50
+        calls = [("readline",), ("read", 5), ("read",), ("getvalue",),
+                 ("truncate", 3), ("write", "ab"), ("__next__",)]
+        for name, *args in calls:
+            for n in range(3):
+                with self.subTest(method=name, n=n):
+                    memio = self.ioclass()
+                    for _ in range(50):
+                        memio.write(piece)
+                    memio.seek(0)
+                    method = getattr(memio, name)
+                    try:
+                        with support.inject_memory_error_cm(n, n + 1):
+                            method(*args)
+                    except MemoryError:
+                        pass
+                    else:
+                        continue
+                    # The contents are unchanged.
+                    self.assertEqual(memio.getvalue(), text)
+                    self.assertEqual(memio.seek(0, 2), len(text))
+                    memio.seek(0)
+                    self.assertEqual(memio.readline(), text)
+                    memio.write("!")
+                    self.assertEqual(memio.getvalue(), text + "!")
+
 
 class CStringIOPickleTest(PyStringIOPickleTest):
     UnsupportedOperation = io.UnsupportedOperation

Misc/NEWS.d/next/Library/2026-10-22-08-15-56.gh-issue-160324.sIoRlz.rst

@@ -0,0 +1,3 @@
+Fix a crash or garbage data when an :class:`io.StringIO` object is used again
+after one of its methods failed with :exc:`MemoryError`.  The object now keeps
+its contents.

Modules/_io/stringio.c

@@ -1,6 +1,7 @@
 #include "Python.h"
 #include <stddef.h>               // offsetof()
 #include "pycore_object.h"
+#include "pycore_unicodeobject.h" // _PyUnicode_ResizeCompact()
 #include "pycore_weakref.h"       // FT_CLEAR_WEAKREFS()
 #include "_iomodule.h"
 
@@ -143,25 +144,52 @@ resize_buffer(stringio *self, size_t size)
     return -1;
 }
 
+/* Shrink the buffer of the writer to its contents, so that the following
+   PyUnicodeWriter_Finish() has nothing to allocate and cannot fail.
+   If this fails, the writer is unchanged. */
+static int
+shrink_writer(stringio *self)
+{
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter *)self->writer;
+    if (writer->readonly || writer->buffer == NULL
+        || writer->pos == 0 || writer->pos == writer->size)
+    {
+        return 0;
+    }
+    PyObject *buffer = _PyUnicode_ResizeCompact(writer->buffer, writer->pos);
+    if (buffer == NULL) {
+        return -1;
+    }
+    writer->buffer = NULL;
+    _PyUnicodeWriter_SetBuffer(writer, buffer);
+    return 0;
+}
+
 static PyObject *
 make_intermediate(stringio *self)
 {
-    PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
-    self->writer = NULL;
-    self->state = STATE_REALIZED;
-    if (intermediate == NULL)
+    /* Create the new writer and shrink the old one first: if this fails,
+       the object keeps its contents. */
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
+    if (writer == NULL)
         return NULL;
+    if (shrink_writer(self) < 0) {
+        PyUnicodeWriter_Discard(writer);
+        return NULL;
+    }
 
-    self->writer = PyUnicodeWriter_Create(0);
-    if (self->writer == NULL) {
-        Py_DECREF(intermediate);
+    PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
+    self->writer = writer;
+    if (intermediate == NULL) {
+        /* The contents are lost, the new writer is empty. */
+        self->string_size = 0;
         return NULL;
     }
     if (PyUnicodeWriter_WriteStr(self->writer, intermediate)) {
         Py_DECREF(intermediate);
+        self->string_size = 0;
         return NULL;
     }
-    self->state = STATE_ACCUMULATING;
     return intermediate;
 }
 
@@ -174,23 +202,30 @@ realize(stringio *self)
     if (self->state == STATE_REALIZED)
         return 0;
     assert(self->state == STATE_ACCUMULATING);
+
+    /* Allocate the internal buffer and shrink the writer first: if this
+       fails, the object stays in the accumulating state and keeps its
+       contents. */
+    if (resize_buffer(self, self->string_size) < 0)
+        return -1;
+    if (shrink_writer(self) < 0)
+        return -1;
     self->state = STATE_REALIZED;
 
     intermediate = PyUnicodeWriter_Finish(self->writer);
     self->writer = NULL;
-    if (intermediate == NULL)
+    if (intermediate == NULL) {
+        /* The contents are lost, the internal buffer contains nothing. */
+        self->string_size = 0;
         return -1;
+    }
 
-    /* Append the intermediate string to the internal buffer.
-       The length should be equal to the current cursor position.
-     */
+    /* Copy the intermediate string to the internal buffer. */
     len = PyUnicode_GET_LENGTH(intermediate);
-    if (resize_buffer(self, len) < 0) {
-        Py_DECREF(intermediate);
-        return -1;
-    }
+    assert(len == self->string_size);
     if (!PyUnicode_AsUCS4(intermediate, self->buf, len, 0)) {
         Py_DECREF(intermediate);
+        self->string_size = 0;
         return -1;
     }
 
@@ -358,7 +393,9 @@ _io_StringIO_read_impl(stringio *self, Py_ssize_t size)
     /* Optimization for seek(0); read() */
     if (self->state == STATE_ACCUMULATING && self->pos == 0 && size == n) {
         PyObject *result = make_intermediate(self);
-        self->pos = self->string_size;
+        if (result != NULL) {
+            self->pos = self->string_size;
+        }
         return result;
     }