streams: Optimize stream.write for bytes & bytearrays

Author: 1st1

tests/test_tcp.py

@@ -24,7 +24,9 @@ async def handle_client(reader, writer):
 
             data = await reader.readexactly(len(B_DATA))
             self.assertEqual(data, B_DATA)
-            writer.writelines([b'SP', bytearray(b'A'), memoryview(b'M')])
+            writer.writelines([b'S', b'P'])
+            writer.write(bytearray(b'A'))
+            writer.write(memoryview(b'M'))
 
             await writer.drain()
             writer.close()

uvloop/handles/stream.pyx

@@ -12,6 +12,9 @@ cdef class _StreamWriteContext:
         uv.uv_buf_t     uv_buf
         Py_buffer       py_buf
 
+        bint            used_buf
+        object          data
+
         UVStream        stream
 
         bint            closed
@@ -21,8 +24,10 @@ cdef class _StreamWriteContext:
             return
 
         self.closed = 1
-        PyBuffer_Release(&self.py_buf)  # void
+        if self.used_buf:
+            PyBuffer_Release(&self.py_buf)  # void
         self.req.data = NULL
+        self.data = None
         IF not DEBUG:
             self.stream = None
         Py_DECREF(self)
@@ -33,11 +38,26 @@ cdef class _StreamWriteContext:
         ctx = _StreamWriteContext.__new__(_StreamWriteContext)
         ctx.stream = None
         ctx.closed = 1
+        ctx.used_buf = 0
 
         ctx.req.data = <void*> ctx
 
-        PyObject_GetBuffer(data, &ctx.py_buf, PyBUF_SIMPLE)
-        ctx.uv_buf = uv.uv_buf_init(<char*>ctx.py_buf.buf, ctx.py_buf.len)
+        if PyBytes_CheckExact(data):
+            ctx.data = data
+            ctx.uv_buf.base = PyBytes_AS_STRING(data)
+            ctx.uv_buf.len = Py_SIZE(data)
+
+        elif PyByteArray_CheckExact(data):
+            ctx.data = data
+            ctx.uv_buf.base = PyByteArray_AS_STRING(data)
+            ctx.uv_buf.len = Py_SIZE(data)
+
+        else:
+            PyObject_GetBuffer(data, &ctx.py_buf, PyBUF_SIMPLE)
+            ctx.used_buf = 1
+            ctx.uv_buf.base = <char*>ctx.py_buf.buf
+            ctx.uv_buf.len = ctx.py_buf.len
+
         ctx.stream = stream
 
         IF DEBUG:

uvloop/includes/python.pxd

@@ -4,5 +4,8 @@ cdef extern from "Python.h":
     void* PyMem_Calloc(size_t nelem, size_t elsize)  # Python >= 3.5!
     void PyMem_Free(void *p)
 
+    char* PyByteArray_AS_STRING(object o)
+    int PyByteArray_CheckExact(object o)
+
     object PyUnicode_EncodeFSDefault(object)
     void PyErr_SetInterrupt() nogil

uvloop/loop.pyx

@@ -11,7 +11,9 @@ from .includes cimport system
 from .includes.python cimport PyMem_Malloc, PyMem_Free, \
                               PyMem_Calloc, PyMem_Realloc, \
                               PyUnicode_EncodeFSDefault, \
-                              PyErr_SetInterrupt
+                              PyErr_SetInterrupt, \
+                              PyByteArray_AS_STRING, \
+                              PyByteArray_CheckExact
 
 from libc.stdint cimport uint64_t
 from libc.string cimport memset, strerror
@@ -22,7 +24,9 @@ from cpython cimport PyErr_CheckSignals, PyErr_Occurred
 from cpython cimport PyThread_get_thread_ident
 from cpython cimport Py_INCREF, Py_DECREF, Py_XDECREF, Py_XINCREF
 from cpython cimport PyObject_GetBuffer, PyBuffer_Release, PyBUF_SIMPLE, \
-                     Py_buffer, PyBytes_AsString
+                     Py_buffer, PyBytes_AsString, PyBytes_CheckExact, \
+                     Py_SIZE, PyBytes_AS_STRING
+
 from cpython cimport PyErr_CheckSignals