Add read_multiple_bytes/read_multiple_bytes_packed methods

capnp/lib/capnp.pyx

@@ -2853,6 +2853,36 @@ class _StructModule(object):
         :rtype: Iterable with elements of :class:`_DynamicStructReader`"""
         reader = _MultiplePackedMessageReader(file.fileno(), self.schema, traversal_limit_in_words, nesting_limit)
         return reader
+    def read_multiple_bytes(self, buf, traversal_limit_in_words = None, nesting_limit = None):
+        """Returns an iterable, that when traversed will return Readers for messages.
+
+        :type buf: buffer
+        :param buf: Any Python object that supports the buffer interface.
+
+        :type traversal_limit_in_words: int
+        :param traversal_limit_in_words: Limits how many total words of data are allowed to be traversed. Is actually a uint64_t, and values can be up to 2^64-1. Default is 8*1024*1024.
+
+        :type nesting_limit: int
+        :param nesting_limit: Limits how many total words of data are allowed to be traversed. Default is 64.
+
+        :rtype: Iterable with elements of :class:`_DynamicStructReader`"""
+        reader = _MultipleBytesMessageReader(buf, self.schema, traversal_limit_in_words, nesting_limit)
+        return reader
+    def read_multiple_bytes_packed(self, buf, traversal_limit_in_words = None, nesting_limit = None):
+        """Returns an iterable, that when traversed will return Readers for messages.
+
+        :type buf: buffer
+        :param buf: Any Python object that supports the buffer interface.
+
+        :type traversal_limit_in_words: int
+        :param traversal_limit_in_words: Limits how many total words of data are allowed to be traversed. Is actually a uint64_t, and values can be up to 2^64-1. Default is 8*1024*1024.
+
+        :type nesting_limit: int
+        :param nesting_limit: Limits how many total words of data are allowed to be traversed. Default is 64.
+
+        :rtype: Iterable with elements of :class:`_DynamicStructReader`"""
+        reader = _MultipleBytesPackedMessageReader(buf, self.schema, traversal_limit_in_words, nesting_limit)
+        return reader
     def from_bytes(self, buf, traversal_limit_in_words = None, nesting_limit = None, builder=False):
         """Returns a Reader for the unpacked object in buf.
 
@@ -3426,6 +3456,78 @@ cdef class _MultiplePackedMessageReader:
     def __iter__(self):
         return self
 
+cdef class _MultipleBytesMessageReader:
+    cdef schema_cpp.ArrayInputStream * stream
+    cdef schema_cpp.BufferedInputStream * buffered_stream
+
+    cdef public object traversal_limit_in_words, nesting_limit, schema, buf
+
+    def __init__(self, buf, schema, traversal_limit_in_words = None, nesting_limit = None):
+        self.schema = schema
+        self.traversal_limit_in_words = traversal_limit_in_words
+        self.nesting_limit = nesting_limit
+
+        cdef const void *ptr
+        cdef Py_ssize_t sz
+        PyObject_AsReadBuffer(buf, &ptr, &sz)
+
+        self.buf = buf
+        self.stream = new schema_cpp.ArrayInputStream(schema_cpp.ByteArrayPtr(<byte *>ptr, sz))
+        self.buffered_stream = new schema_cpp.BufferedInputStreamWrapper(deref(self.stream))
+
+    def __dealloc__(self):
+        del self.stream
+        del self.buffered_stream
+
+    def __next__(self):
+        try:
+            reader = _InputMessageReader()._init(deref(self.buffered_stream), self.traversal_limit_in_words, self.nesting_limit, self)
+            return reader.get_root(self.schema)
+        except KjException as e:
+            if 'EOF' in str(e):
+                raise StopIteration
+            else:
+                raise
+
+    def __iter__(self):
+        return self
+
+cdef class _MultipleBytesPackedMessageReader:
+    cdef schema_cpp.ArrayInputStream * stream
+    cdef schema_cpp.BufferedInputStream * buffered_stream
+
+    cdef public object traversal_limit_in_words, nesting_limit, schema, buf
+
+    def __init__(self, buf, schema, traversal_limit_in_words = None, nesting_limit = None):
+        self.schema = schema
+        self.traversal_limit_in_words = traversal_limit_in_words
+        self.nesting_limit = nesting_limit
+
+        cdef const void *ptr
+        cdef Py_ssize_t sz
+        PyObject_AsReadBuffer(buf, &ptr, &sz)
+
+        self.buf = buf
+        self.stream = new schema_cpp.ArrayInputStream(schema_cpp.ByteArrayPtr(<byte *>ptr, sz))
+        self.buffered_stream = new schema_cpp.BufferedInputStreamWrapper(deref(self.stream))
+
+    def __dealloc__(self):
+        del self.stream
+        del self.buffered_stream
+
+    def __next__(self):
+        try:
+            reader = _PackedMessageReader()._init(deref(self.buffered_stream), self.traversal_limit_in_words, self.nesting_limit, self)
+            return reader.get_root(self.schema)
+        except KjException as e:
+            if 'EOF' in str(e):
+                raise StopIteration
+            else:
+                raise
+
+    def __iter__(self):
+        return self
+
 @cython.internal
 cdef class _AlignedBuffer:
     cdef char * buf

test/test_serialization.py

@@ -60,6 +60,17 @@ def test_roundtrip_file_multiple(all_types):
     for msg in all_types.TestAllTypes.read_multiple(f):
         test_regression.check_all_types(msg)
 
+def test_roundtrip_bytes_multiple(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    msgs = msg.to_bytes()
+    msgs += msg.to_bytes()
+    msgs += msg.to_bytes()
+
+    for msg in all_types.TestAllTypes.read_multiple_bytes(msgs):
+        test_regression.check_all_types(msg)
+
 def test_roundtrip_file_multiple_packed(all_types):
     f = tempfile.TemporaryFile()
     msg = all_types.TestAllTypes.new_message()
@@ -72,6 +83,17 @@ def test_roundtrip_file_multiple_packed(all_types):
     for msg in all_types.TestAllTypes.read_multiple_packed(f):
         test_regression.check_all_types(msg)
 
+def test_roundtrip_bytes_multiple_packed(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    msgs = msg.to_bytes_packed()
+    msgs += msg.to_bytes_packed()
+    msgs += msg.to_bytes_packed()
+
+    for msg in all_types.TestAllTypes.read_multiple_bytes_packed(msgs):
+        test_regression.check_all_types(msg)
+
 def test_roundtrip_dict(all_types):
     msg = all_types.TestAllTypes.new_message()
     test_regression.init_all_types(msg)