80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import pytest
|
|
import capnp # noqa: F401
|
|
import os
|
|
|
|
this_dir = os.path.dirname(__file__)
|
|
|
|
|
|
@pytest.fixture
|
|
def all_types():
|
|
return capnp.load(os.path.join(this_dir, "all_types.capnp"))
|
|
|
|
|
|
def test_bytearray_allocator(all_types):
|
|
class Allocator:
|
|
def __init__(self):
|
|
self.cur_size = 0
|
|
self.last_size = 0
|
|
|
|
def __call__(self, minimum_size: int) -> bytearray:
|
|
actual_size = max(minimum_size, self.cur_size)
|
|
self.last_size = actual_size
|
|
self.cur_size += actual_size
|
|
WORD_SIZE = 8
|
|
byte_count = actual_size * WORD_SIZE
|
|
return bytearray(byte_count)
|
|
|
|
allocator = Allocator()
|
|
assert allocator.cur_size == 0
|
|
assert allocator.last_size == 0
|
|
msg_builder = capnp._PyCustomMessageBuilder(allocator, 1024)
|
|
struct_builder = msg_builder.init_root(all_types.TestAllTypes)
|
|
assert allocator.cur_size == 1024
|
|
assert allocator.last_size == 1024
|
|
|
|
struct_builder.init("dataField", 5)
|
|
assert bytes(struct_builder._get("dataField")) == b"\x00\x00\x00\x00\x00"
|
|
|
|
struct_builder.dataField = b"hello"
|
|
assert bytes(struct_builder._get("dataField")) == b"hello"
|
|
|
|
struct_reader = struct_builder.as_reader()
|
|
assert bytes(struct_reader._get("dataField")) == b"hello"
|
|
|
|
|
|
def test_memoryview_allocator(all_types):
|
|
class MemoryViewAllocator:
|
|
def __init__(self):
|
|
self.cur_size = 0
|
|
self.last_size = 0
|
|
self.buffers = []
|
|
|
|
def __call__(self, minimum_size: int) -> memoryview:
|
|
actual_size = max(minimum_size, self.cur_size)
|
|
self.last_size = actual_size
|
|
self.cur_size += actual_size
|
|
WORD_SIZE = 8
|
|
byte_count = actual_size * WORD_SIZE
|
|
buffer = bytearray(byte_count)
|
|
self.buffers.append(buffer)
|
|
return memoryview(buffer)
|
|
|
|
allocator = MemoryViewAllocator()
|
|
assert allocator.cur_size == 0
|
|
assert allocator.last_size == 0
|
|
msg_builder = capnp._PyCustomMessageBuilder(allocator, 1024)
|
|
struct_builder = msg_builder.init_root(all_types.TestAllTypes)
|
|
assert allocator.cur_size == 1024
|
|
assert allocator.last_size == 1024
|
|
|
|
struct_builder.init("dataField", 5)
|
|
assert bytes(struct_builder._get("dataField")) == b"\x00\x00\x00\x00\x00"
|
|
|
|
struct_builder.dataField = b"hello"
|
|
assert bytes(struct_builder._get("dataField")) == b"hello"
|
|
|
|
struct_reader = struct_builder.as_reader()
|
|
assert bytes(struct_reader._get("dataField")) == b"hello"
|