Optimize numeric array conversion and serialize directly into Python bytes

capnp/lib/capnp.pyx

@@ -13,6 +13,7 @@ from capnp.helpers.helpers cimport init_capnp_api
 from builtins import memoryview as BuiltinsMemoryview
 from cpython cimport Py_buffer, PyObject_CheckBuffer
 from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_CONTIG_RO
+from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING
 from cpython.exc cimport PyErr_Clear
 from cython.operator cimport dereference as deref
 from libc.stdlib cimport malloc, free
@@ -28,6 +29,11 @@ import warnings as _warnings
 
 from types import ModuleType as _ModuleType
 
+cdef object _numpy = None
+
+cdef extern from "capnp/helpers/serialize.h":
+    object messageToPythonBytes(schema_cpp.MessageBuilder&) except +reraise_kj_exception
+
 _CAPNP_VERSION_MAJOR = capnp.CAPNP_VERSION_MAJOR
 _CAPNP_VERSION_MINOR = capnp.CAPNP_VERSION_MINOR
 _CAPNP_VERSION_MICRO = capnp.CAPNP_VERSION_MICRO
@@ -243,6 +249,49 @@ cdef class _DynamicListReader:
         ptr = self.thisptr[index]
         return to_python_reader(ptr, self._parent)
 
+    def to_float64_bytes(self):
+        cdef Py_ssize_t size = self.thisptr.size()
+        cdef bytes result = PyBytes_FromStringAndSize(NULL, size * sizeof(double))
+        cdef char* target = PyBytes_AS_STRING(result)
+        cdef Py_ssize_t index
+        cdef double value
+        cdef C_DynamicValue.Reader item
+        cdef int kind
+        if size:
+            item = self.thisptr[0]
+            kind = item.getType()
+            if kind != capnp.TYPE_FLOAT and kind != capnp.TYPE_INT and kind != capnp.TYPE_UINT:
+                raise TypeError("to_float64_bytes requires a numeric list")
+        for index in range(size):
+            item = self.thisptr[index]
+            value = item.asDouble()
+            memcpy(target + index * sizeof(double), &value, sizeof(double))
+        return result
+
+    def __array__(self, dtype=None, copy=None):
+        global _numpy
+        if _numpy is None:
+            _numpy = __import__("numpy")
+        np = _numpy
+        if copy is False:
+            raise ValueError("Cap'n Proto lists require a copy for NumPy conversion")
+        cdef Py_ssize_t size = self.thisptr.size()
+        cdef Py_ssize_t index
+        cdef C_DynamicValue.Reader item
+        cdef int kind
+        cdef double[::1] values64
+        if size:
+            item = self.thisptr[0]
+            kind = item.getType()
+            if kind != capnp.TYPE_FLOAT:
+                return np.array(list(self), dtype=dtype)
+        result = np.empty(size, dtype=np.float64)
+        values64 = result
+        for index in range(size):
+            item = self.thisptr[index]
+            values64[index] = item.asDouble()
+        return result if dtype is None else result.astype(dtype, copy=False)
+
     def __getitem__(self, int64_t index):
         cdef uint size = self.thisptr.size()
         if index >= size:
@@ -289,6 +338,9 @@ cdef class _DynamicListBuilder:
         ptr = self.thisptr[index]
         return to_python_builder(ptr, self._parent)
 
+    def __array__(self, dtype=None, copy=None):
+        return _DynamicListReader()._init(self.thisptr.asReader(), self._parent).__array__(dtype=dtype, copy=copy)
+
     def __getitem__(self, int64_t index):
         cdef uint size = self.thisptr.size()
         if index >= size:
@@ -888,17 +940,13 @@ cdef class _DynamicStructBuilder:
     cpdef to_bytes(_DynamicStructBuilder self):
         """Returns the struct's containing message as a Python bytes object in the unpacked binary format.
 
-        This is inefficient; it makes several copies.
-
         :rtype: bytes
 
         :Raises: :exc:`KjException` if this isn't the message's root struct.
         """
         self._check_write()
         cdef _MessageBuilder builder = self._parent
-        array = schema_cpp.messageToFlatArray(deref(builder.thisptr))
-        cdef const char* ptr = <const char *>array.begin()
-        cdef bytes ret = ptr[:8*array.size()]
+        cdef bytes ret = messageToPythonBytes(deref(builder.thisptr))
         self._is_written = True
         return ret