diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 2f76d3f..b3baf89 100644 --- a/capnp/lib/capnp.pyx +++ b/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(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(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 diff --git a/test/test_serialization.py b/test/test_serialization.py index 3928c31..d1bfc66 100644 --- a/test/test_serialization.py +++ b/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)