From 144acbea7a4e6fb13e9307fd481714fcc57423d6 Mon Sep 17 00:00:00 2001 From: Yevhenii Kizim Date: Mon, 21 Oct 2024 16:43:14 +0200 Subject: [PATCH] Add structure-free read_multiple_bytes_packed Motivation: A server sends data packages that consist of multiple serialized capnproto messages of different structures. Every message is guaranteed to have the same first field, which works as a message header containing information about the message structure type. The scheme comprises the `UnknownMessage` structure that allows parsing the header only. Solution: provide a public interface that iterates buffer with AnyPointer readers to cast a message to `UnknownMessage` first and then to a specific structure type. --- capnp/lib/capnp.pyx | 57 +++++++++++++++++++++ test/test_structs_sequence.py | 93 +++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 test/test_structs_sequence.py diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 5ec51b1..34fc63c 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -4119,6 +4119,44 @@ cdef class _MultipleBytesPackedMessageReader: return self +cdef class _MultipleBytesPackedAnyMessageReader: + cdef schema_cpp.ArrayInputStream * stream + cdef schema_cpp.BufferedInputStream * buffered_stream + cdef Py_buffer view + + cdef public object traversal_limit_in_words, nesting_limit, schema, buf + + def __init__(self, buf, traversal_limit_in_words=None, nesting_limit=None): + self.traversal_limit_in_words = traversal_limit_in_words + self.nesting_limit = nesting_limit + + if PyObject_GetBuffer(buf, &self.view, PyBUF_SIMPLE) != 0: + raise KjException("could not get read buffer") + + self.buf = buf + self.stream = new schema_cpp.ArrayInputStream(schema_cpp.ByteArrayPtr(self.view.buf, self.view.len)) + self.buffered_stream = new schema_cpp.BufferedInputStreamWrapper(deref(self.stream)) + + def __dealloc__(self): + PyBuffer_Release(&self.view) + del self.buffered_stream + del self.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_as_any() + 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 @@ -4418,6 +4456,25 @@ def load(file_name, display_name=None, imports=[]): return _global_schema_parser.load(file_name, display_name, imports) +def read_multiple_bytes_packed(buf, traversal_limit_in_words=None, nesting_limit=None): + """Returns an iterable, that when traversed will return Readers for AnyPointer 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 = _MultipleBytesPackedAnyMessageReader(buf, traversal_limit_in_words, nesting_limit) + return reader + + # Automatically include the system and built-in capnp paths # Highest priority at position 0 _capnp_paths = [ diff --git a/test/test_structs_sequence.py b/test/test_structs_sequence.py new file mode 100644 index 0000000..e2b5ed4 --- /dev/null +++ b/test/test_structs_sequence.py @@ -0,0 +1,93 @@ +import os + +import pytest + +import capnp + +this_dir = os.path.dirname(__file__) + + +@pytest.fixture +def message_schemas(): + return capnp.load(os.path.join(this_dir, "test_structs_sequence.capnp")) + + +@pytest.fixture +def make_apple(message_schemas): + def _make_apple(color: str): + apple = message_schemas.Apple.new_message() + apple.fruitId = message_schemas.FruitId.apple + apple.color = color + return apple + + return _make_apple + + +@pytest.fixture +def red_apple(make_apple): + return make_apple("Red") + + +@pytest.fixture +def green_apple(make_apple): + return make_apple("Green") + + +@pytest.fixture +def banana(message_schemas): + banana_ = message_schemas.Banana.new_message() + banana_.fruitId = message_schemas.FruitId.banana + banana_.length = 12.345 + return banana_ + + +@pytest.fixture +def cherry(message_schemas): + cherry_ = message_schemas.Cherry.new_message() + cherry_.fruitId = message_schemas.FruitId.cherry + cherry_.sweetness = 64 + return cherry_ + + +@pytest.fixture +def fruit_basket(cherry, red_apple, banana, green_apple): + return [cherry, red_apple, banana, green_apple] + + +@pytest.fixture +def fruit_basket_encoded(fruit_basket): + return b"".join(fruit.to_bytes_packed() for fruit in fruit_basket) + + +@pytest.fixture +def expected(fruit_basket): + return [fruit.to_dict() for fruit in fruit_basket] + + +def test_parse_structs_sequence(message_schemas, fruit_basket_encoded, expected): + # ARRANGE + reader = capnp.read_multiple_bytes_packed(fruit_basket_encoded) + + def _parse_fruit(any_): + unknown_fruit = any_.as_struct(message_schemas.UnknownFruit) + if unknown_fruit.fruitId == message_schemas.FruitId.apple: + return any_.as_struct(message_schemas.Apple) + + if unknown_fruit.fruitId == message_schemas.FruitId.banana: + return any_.as_struct(message_schemas.Banana) + + if unknown_fruit.fruitId == message_schemas.FruitId.cherry: + return any_.as_struct(message_schemas.Cherry) + + return unknown_fruit + + # ACT + parsed = [_parse_fruit(any_).to_dict() for any_ in reader] + + # ASSERT + assert parsed == expected + + +def test_empty_sequence(): + reader = capnp.read_multiple_bytes_packed(b"") + assert len(list(reader)) == 0