Skip to content

Commit bfba8c3

Browse files
authored
bpo-36748: optimize TextIOWrapper.write() for ASCII string (GH-13002)
1 parent 8a533ff commit bfba8c3

File tree

2 files changed

+98
-21
lines changed

2 files changed

+98
-21
lines changed
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Optimized write buffering in C implementation of ``TextIOWrapper``. Writing
2+
ASCII string to ``TextIOWrapper`` with ascii, latin1, or utf-8 encoding is
3+
about 20% faster. Patch by Inada Naoki.

Modules/_io/textio.c

Lines changed: 95 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -674,8 +674,8 @@ typedef struct
674674
*/
675675
PyObject *decoded_chars; /* buffer for text returned from decoder */
676676
Py_ssize_t decoded_chars_used; /* offset into _decoded_chars for read() */
677-
PyObject *pending_bytes; /* list of bytes objects waiting to be
678-
written, or NULL */
677+
PyObject *pending_bytes; // data waiting to be written.
678+
// ascii unicode, bytes, or list of them.
679679
Py_ssize_t pending_bytes_count;
680680

681681
/* snapshot is either NULL, or a tuple (dec_flags, next_input) where
@@ -777,6 +777,15 @@ latin1_encode(textio *self, PyObject *text)
777777
return _PyUnicode_AsLatin1String(text, PyUnicode_AsUTF8(self->errors));
778778
}
779779

780+
// Return true when encoding can be skipped when text is ascii.
781+
static inline int
782+
is_asciicompat_encoding(encodefunc_t f)
783+
{
784+
return f == (encodefunc_t) ascii_encode
785+
|| f == (encodefunc_t) latin1_encode
786+
|| f == (encodefunc_t) utf8_encode;
787+
}
788+
780789
/* Map normalized encoding names onto the specialized encoding funcs */
781790

782791
typedef struct {
@@ -1489,21 +1498,62 @@ _io_TextIOWrapper_detach_impl(textio *self)
14891498
static int
14901499
_textiowrapper_writeflush(textio *self)
14911500
{
1492-
PyObject *pending, *b, *ret;
1493-
14941501
if (self->pending_bytes == NULL)
14951502
return 0;
14961503

1497-
pending = self->pending_bytes;
1498-
Py_INCREF(pending);
1499-
self->pending_bytes_count = 0;
1500-
Py_CLEAR(self->pending_bytes);
1504+
PyObject *pending = self->pending_bytes;
1505+
PyObject *b;
15011506

1502-
b = _PyBytes_Join(_PyIO_empty_bytes, pending);
1507+
if (PyBytes_Check(pending)) {
1508+
b = pending;
1509+
Py_INCREF(b);
1510+
}
1511+
else if (PyUnicode_Check(pending)) {
1512+
assert(PyUnicode_IS_ASCII(pending));
1513+
assert(PyUnicode_GET_LENGTH(pending) == self->pending_bytes_count);
1514+
b = PyBytes_FromStringAndSize(
1515+
PyUnicode_DATA(pending), PyUnicode_GET_LENGTH(pending));
1516+
if (b == NULL) {
1517+
return -1;
1518+
}
1519+
}
1520+
else {
1521+
assert(PyList_Check(pending));
1522+
b = PyBytes_FromStringAndSize(NULL, self->pending_bytes_count);
1523+
if (b == NULL) {
1524+
return -1;
1525+
}
1526+
1527+
char *buf = PyBytes_AsString(b);
1528+
Py_ssize_t pos = 0;
1529+
1530+
for (Py_ssize_t i = 0; i < PyList_GET_SIZE(pending); i++) {
1531+
PyObject *obj = PyList_GET_ITEM(pending, i);
1532+
char *src;
1533+
Py_ssize_t len;
1534+
if (PyUnicode_Check(obj)) {
1535+
assert(PyUnicode_IS_ASCII(obj));
1536+
src = PyUnicode_DATA(obj);
1537+
len = PyUnicode_GET_LENGTH(obj);
1538+
}
1539+
else {
1540+
assert(PyBytes_Check(obj));
1541+
if (PyBytes_AsStringAndSize(obj, &src, &len) < 0) {
1542+
Py_DECREF(b);
1543+
return -1;
1544+
}
1545+
}
1546+
memcpy(buf + pos, src, len);
1547+
pos += len;
1548+
}
1549+
assert(pos == self->pending_bytes_count);
1550+
}
1551+
1552+
self->pending_bytes_count = 0;
1553+
self->pending_bytes = NULL;
15031554
Py_DECREF(pending);
1504-
if (b == NULL)
1505-
return -1;
1506-
ret = NULL;
1555+
1556+
PyObject *ret;
15071557
do {
15081558
ret = PyObject_CallMethodObjArgs(self->buffer,
15091559
_PyIO_str_write, b, NULL);
@@ -1566,37 +1616,61 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
15661616

15671617
/* XXX What if we were just reading? */
15681618
if (self->encodefunc != NULL) {
1569-
b = (*self->encodefunc)((PyObject *) self, text);
1619+
if (PyUnicode_IS_ASCII(text) && is_asciicompat_encoding(self->encodefunc)) {
1620+
b = text;
1621+
Py_INCREF(b);
1622+
}
1623+
else {
1624+
b = (*self->encodefunc)((PyObject *) self, text);
1625+
}
15701626
self->encoding_start_of_stream = 0;
15711627
}
15721628
else
15731629
b = PyObject_CallMethodObjArgs(self->encoder,
15741630
_PyIO_str_encode, text, NULL);
1631+
15751632
Py_DECREF(text);
15761633
if (b == NULL)
15771634
return NULL;
1578-
if (!PyBytes_Check(b)) {
1635+
if (b != text && !PyBytes_Check(b)) {
15791636
PyErr_Format(PyExc_TypeError,
15801637
"encoder should return a bytes object, not '%.200s'",
15811638
Py_TYPE(b)->tp_name);
15821639
Py_DECREF(b);
15831640
return NULL;
15841641
}
15851642

1643+
Py_ssize_t bytes_len;
1644+
if (b == text) {
1645+
bytes_len = PyUnicode_GET_LENGTH(b);
1646+
}
1647+
else {
1648+
bytes_len = PyBytes_GET_SIZE(b);
1649+
}
1650+
15861651
if (self->pending_bytes == NULL) {
1587-
self->pending_bytes = PyList_New(0);
1588-
if (self->pending_bytes == NULL) {
1652+
self->pending_bytes_count = 0;
1653+
self->pending_bytes = b;
1654+
}
1655+
else if (!PyList_CheckExact(self->pending_bytes)) {
1656+
PyObject *list = PyList_New(2);
1657+
if (list == NULL) {
15891658
Py_DECREF(b);
15901659
return NULL;
15911660
}
1592-
self->pending_bytes_count = 0;
1661+
PyList_SET_ITEM(list, 0, self->pending_bytes);
1662+
PyList_SET_ITEM(list, 1, b);
1663+
self->pending_bytes = list;
15931664
}
1594-
if (PyList_Append(self->pending_bytes, b) < 0) {
1665+
else {
1666+
if (PyList_Append(self->pending_bytes, b) < 0) {
1667+
Py_DECREF(b);
1668+
return NULL;
1669+
}
15951670
Py_DECREF(b);
1596-
return NULL;
15971671
}
1598-
self->pending_bytes_count += PyBytes_GET_SIZE(b);
1599-
Py_DECREF(b);
1672+
1673+
self->pending_bytes_count += bytes_len;
16001674
if (self->pending_bytes_count > self->chunk_size || needflush ||
16011675
text_needflush) {
16021676
if (_textiowrapper_writeflush(self) < 0)

0 commit comments

Comments
 (0)