Skip to content

Commit

Permalink
pythongh-125196: Use PyUnicodeWriter in _io.StringIO
Browse files Browse the repository at this point in the history
  • Loading branch information
vstinner committed Oct 15, 2024
1 parent c8a1818 commit 5e96bfb
Showing 1 changed file with 31 additions and 18 deletions.
49 changes: 31 additions & 18 deletions Modules/_io/stringio.c
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ typedef struct {
_PyUnicodeWriter is destroyed.
*/
int state;
_PyUnicodeWriter writer;
PyUnicodeWriter *writer;

char ok; /* initialized? */
char closed;
Expand Down Expand Up @@ -129,14 +129,20 @@ resize_buffer(stringio *self, size_t size)
static PyObject *
make_intermediate(stringio *self)
{
PyObject *intermediate = _PyUnicodeWriter_Finish(&self->writer);
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
self->writer = NULL;
self->state = STATE_REALIZED;
if (intermediate == NULL)
if (intermediate == NULL) {
return NULL;
}

self->writer = PyUnicodeWriter_Create(0);
if (self->writer == NULL) {
Py_DECREF(intermediate);
return NULL;
}

_PyUnicodeWriter_Init(&self->writer);
self->writer.overallocate = 1;
if (_PyUnicodeWriter_WriteStr(&self->writer, intermediate)) {
if (PyUnicodeWriter_WriteStr(self->writer, intermediate)) {
Py_DECREF(intermediate);
return NULL;
}
Expand All @@ -147,22 +153,21 @@ make_intermediate(stringio *self)
static int
realize(stringio *self)
{
Py_ssize_t len;
PyObject *intermediate;

if (self->state == STATE_REALIZED)
return 0;
assert(self->state == STATE_ACCUMULATING);
self->state = STATE_REALIZED;

intermediate = _PyUnicodeWriter_Finish(&self->writer);
if (intermediate == NULL)
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
self->writer = NULL;
if (intermediate == NULL) {
return -1;
}

/* Append the intermediate string to the internal buffer.
The length should be equal to the current cursor position.
*/
len = PyUnicode_GET_LENGTH(intermediate);
Py_ssize_t len = PyUnicode_GET_LENGTH(intermediate);
if (resize_buffer(self, len) < 0) {
Py_DECREF(intermediate);
return -1;
Expand Down Expand Up @@ -217,8 +222,9 @@ write_str(stringio *self, PyObject *obj)

if (self->state == STATE_ACCUMULATING) {
if (self->string_size == self->pos) {
if (_PyUnicodeWriter_WriteStr(&self->writer, decoded))
if (PyUnicodeWriter_WriteStr(self->writer, decoded)) {
goto fail;
}
goto success;
}
if (realize(self))
Expand Down Expand Up @@ -577,7 +583,8 @@ _io_StringIO_close_impl(stringio *self)
/* Free up some memory */
if (resize_buffer(self, 0) < 0)
return NULL;
_PyUnicodeWriter_Dealloc(&self->writer);
PyUnicodeWriter_Discard(self->writer);
self->writer = NULL;
Py_CLEAR(self->readnl);
Py_CLEAR(self->writenl);
Py_CLEAR(self->decoder);
Expand Down Expand Up @@ -615,7 +622,8 @@ stringio_dealloc(stringio *self)
PyMem_Free(self->buf);
self->buf = NULL;
}
_PyUnicodeWriter_Dealloc(&self->writer);
PyUnicodeWriter_Discard(self->writer);
self->writer = NULL;
(void)stringio_clear(self);
if (self->weakreflist != NULL) {
PyObject_ClearWeakRefs((PyObject *) self);
Expand Down Expand Up @@ -699,7 +707,8 @@ _io_StringIO___init___impl(stringio *self, PyObject *value,

self->ok = 0;

_PyUnicodeWriter_Dealloc(&self->writer);
PyUnicodeWriter_Discard(self->writer);
self->writer = NULL;
Py_CLEAR(self->readnl);
Py_CLEAR(self->writenl);
Py_CLEAR(self->decoder);
Expand Down Expand Up @@ -754,8 +763,12 @@ _io_StringIO___init___impl(stringio *self, PyObject *value,
/* Empty stringio object, we can start by accumulating */
if (resize_buffer(self, 0) < 0)
return -1;
_PyUnicodeWriter_Init(&self->writer);
self->writer.overallocate = 1;

self->writer = PyUnicodeWriter_Create(0);
if (self->writer == NULL) {
return -1;
}

self->state = STATE_ACCUMULATING;
}
self->pos = 0;
Expand Down

0 comments on commit 5e96bfb

Please sign in to comment.