diff options
| -rw-r--r-- | Lib/test/test_io.py | 27 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Library/2021-02-20-12-15-29.bpo-43260.6znAas.rst | 2 | ||||
| -rw-r--r-- | Modules/_io/textio.c | 20 |
3 files changed, 46 insertions, 3 deletions
diff --git a/Lib/test/test_io.py b/Lib/test/test_io.py index cc54d0e..3768b62 100644 --- a/Lib/test/test_io.py +++ b/Lib/test/test_io.py @@ -3767,6 +3767,33 @@ class CTextIOWrapperTest(TextIOWrapperTest): with self.assertRaises(AttributeError): del t._CHUNK_SIZE + def test_internal_buffer_size(self): + # bpo-43260: TextIOWrapper's internal buffer should not store + # data larger than chunk size. + chunk_size = 8192 # default chunk size, updated later + + class MockIO(self.MockRawIO): + def write(self, data): + if len(data) > chunk_size: + raise RuntimeError + return super().write(data) + + buf = MockIO() + t = self.TextIOWrapper(buf, encoding="ascii") + chunk_size = t._CHUNK_SIZE + t.write("abc") + t.write("def") + # default chunk size is 8192 bytes so t don't write data to buf. + self.assertEqual([], buf._write_stack) + + with self.assertRaises(RuntimeError): + t.write("x"*(chunk_size+1)) + + self.assertEqual([b"abcdef"], buf._write_stack) + t.write("ghi") + t.write("x"*chunk_size) + self.assertEqual([b"abcdef", b"ghi", b"x"*chunk_size], buf._write_stack) + class PyTextIOWrapperTest(TextIOWrapperTest): io = pyio diff --git a/Misc/NEWS.d/next/Library/2021-02-20-12-15-29.bpo-43260.6znAas.rst b/Misc/NEWS.d/next/Library/2021-02-20-12-15-29.bpo-43260.6znAas.rst new file mode 100644 index 0000000..f3c21d1 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-02-20-12-15-29.bpo-43260.6znAas.rst @@ -0,0 +1,2 @@ +Fix TextIOWrapper can not flush internal buffer forever after very large +text is written. diff --git a/Modules/_io/textio.c b/Modules/_io/textio.c index f08d14e..03001ec 100644 --- a/Modules/_io/textio.c +++ b/Modules/_io/textio.c @@ -1585,6 +1585,8 @@ _textiowrapper_writeflush(textio *self) ret = PyObject_CallMethodOneArg(self->buffer, _PyIO_str_write, b); } while (ret == NULL && _PyIO_trap_eintr()); Py_DECREF(b); + // NOTE: We cleared buffer but we don't know how many bytes are actually written + // when an error occurred. if (ret == NULL) return -1; Py_DECREF(ret); @@ -1642,7 +1644,10 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text) /* XXX What if we were just reading? */ if (self->encodefunc != NULL) { - if (PyUnicode_IS_ASCII(text) && is_asciicompat_encoding(self->encodefunc)) { + if (PyUnicode_IS_ASCII(text) && + // See bpo-43260 + PyUnicode_GET_LENGTH(text) <= self->chunk_size && + is_asciicompat_encoding(self->encodefunc)) { b = text; Py_INCREF(b); } @@ -1651,8 +1656,9 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text) } self->encoding_start_of_stream = 0; } - else + else { b = PyObject_CallMethodOneArg(self->encoder, _PyIO_str_encode, text); + } Py_DECREF(text); if (b == NULL) @@ -1677,6 +1683,14 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text) self->pending_bytes_count = 0; self->pending_bytes = b; } + else if (self->pending_bytes_count + bytes_len > self->chunk_size) { + // Prevent to concatenate more than chunk_size data. + if (_textiowrapper_writeflush(self) < 0) { + Py_DECREF(b); + return NULL; + } + self->pending_bytes = b; + } else if (!PyList_CheckExact(self->pending_bytes)) { PyObject *list = PyList_New(2); if (list == NULL) { @@ -1696,7 +1710,7 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text) } self->pending_bytes_count += bytes_len; - if (self->pending_bytes_count > self->chunk_size || needflush || + if (self->pending_bytes_count >= self->chunk_size || needflush || text_needflush) { if (_textiowrapper_writeflush(self) < 0) return NULL; |