11#include "Python.h"
22#include <stddef.h> // offsetof()
33#include "pycore_object.h"
4+ #include "pycore_unicodeobject.h" // _PyUnicode_ResizeCompact()
45#include "pycore_weakref.h" // FT_CLEAR_WEAKREFS()
56#include "_iomodule.h"
67
@@ -143,25 +144,52 @@ resize_buffer(stringio *self, size_t size)
143144 return -1 ;
144145}
145146
147+ /* Shrink the buffer of the writer to its contents, so that the following
148+ PyUnicodeWriter_Finish() has nothing to allocate and cannot fail.
149+ If this fails, the writer is unchanged. */
150+ static int
151+ shrink_writer (stringio * self )
152+ {
153+ _PyUnicodeWriter * writer = (_PyUnicodeWriter * )self -> writer ;
154+ if (writer -> readonly || writer -> buffer == NULL
155+ || writer -> pos == 0 || writer -> pos == writer -> size )
156+ {
157+ return 0 ;
158+ }
159+ PyObject * buffer = _PyUnicode_ResizeCompact (writer -> buffer , writer -> pos );
160+ if (buffer == NULL ) {
161+ return -1 ;
162+ }
163+ writer -> buffer = NULL ;
164+ _PyUnicodeWriter_SetBuffer (writer , buffer );
165+ return 0 ;
166+ }
167+
146168static PyObject *
147169make_intermediate (stringio * self )
148170{
149- PyObject * intermediate = PyUnicodeWriter_Finish ( self -> writer );
150- self -> writer = NULL ;
151- self -> state = STATE_REALIZED ;
152- if (intermediate == NULL )
171+ /* Create the new writer and shrink the old one first: if this fails,
172+ the object keeps its contents. */
173+ PyUnicodeWriter * writer = PyUnicodeWriter_Create ( 0 ) ;
174+ if (writer == NULL )
153175 return NULL ;
176+ if (shrink_writer (self ) < 0 ) {
177+ PyUnicodeWriter_Discard (writer );
178+ return NULL ;
179+ }
154180
155- self -> writer = PyUnicodeWriter_Create (0 );
156- if (self -> writer == NULL ) {
157- Py_DECREF (intermediate );
181+ PyObject * intermediate = PyUnicodeWriter_Finish (self -> writer );
182+ self -> writer = writer ;
183+ if (intermediate == NULL ) {
184+ /* The contents are lost, the new writer is empty. */
185+ self -> string_size = 0 ;
158186 return NULL ;
159187 }
160188 if (PyUnicodeWriter_WriteStr (self -> writer , intermediate )) {
161189 Py_DECREF (intermediate );
190+ self -> string_size = 0 ;
162191 return NULL ;
163192 }
164- self -> state = STATE_ACCUMULATING ;
165193 return intermediate ;
166194}
167195
@@ -174,23 +202,30 @@ realize(stringio *self)
174202 if (self -> state == STATE_REALIZED )
175203 return 0 ;
176204 assert (self -> state == STATE_ACCUMULATING );
205+
206+ /* Allocate the internal buffer and shrink the writer first: if this
207+ fails, the object stays in the accumulating state and keeps its
208+ contents. */
209+ if (resize_buffer (self , self -> string_size ) < 0 )
210+ return -1 ;
211+ if (shrink_writer (self ) < 0 )
212+ return -1 ;
177213 self -> state = STATE_REALIZED ;
178214
179215 intermediate = PyUnicodeWriter_Finish (self -> writer );
180216 self -> writer = NULL ;
181- if (intermediate == NULL )
217+ if (intermediate == NULL ) {
218+ /* The contents are lost, the internal buffer contains nothing. */
219+ self -> string_size = 0 ;
182220 return -1 ;
221+ }
183222
184- /* Append the intermediate string to the internal buffer.
185- The length should be equal to the current cursor position.
186- */
223+ /* Copy the intermediate string to the internal buffer. */
187224 len = PyUnicode_GET_LENGTH (intermediate );
188- if (resize_buffer (self , len ) < 0 ) {
189- Py_DECREF (intermediate );
190- return -1 ;
191- }
225+ assert (len == self -> string_size );
192226 if (!PyUnicode_AsUCS4 (intermediate , self -> buf , len , 0 )) {
193227 Py_DECREF (intermediate );
228+ self -> string_size = 0 ;
194229 return -1 ;
195230 }
196231
@@ -358,7 +393,9 @@ _io_StringIO_read_impl(stringio *self, Py_ssize_t size)
358393 /* Optimization for seek(0); read() */
359394 if (self -> state == STATE_ACCUMULATING && self -> pos == 0 && size == n ) {
360395 PyObject * result = make_intermediate (self );
361- self -> pos = self -> string_size ;
396+ if (result != NULL ) {
397+ self -> pos = self -> string_size ;
398+ }
362399 return result ;
363400 }
364401
0 commit comments