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.
This commit is contained in:
committed by
Jacob Alexander
parent
237fa7dd4b
commit
144acbea7a
@@ -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(<byte *>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 = [
|
||||
|
||||
93
test/test_structs_sequence.py
Normal file
93
test/test_structs_sequence.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user