From b68bb01a53649ee2063be4dc930566bb62064208 Mon Sep 17 00:00:00 2001 From: Brian Xu Date: Fri, 30 Jan 2026 00:04:49 +1100 Subject: [PATCH 1/3] refine document for py custom message builder --- capnp/lib/capnp.pyx | 57 +++++++++++++++++++++----- examples/py_custom_message_builder.py | 47 +++++++++++++++++++++ test/test_py_custom_message_builder.py | 49 ++++++++++++++++++---- 3 files changed, 134 insertions(+), 19 deletions(-) diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 6449a9940..2439f1a8a 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -1314,10 +1314,11 @@ cdef class _DynamicStructReader: :type num_first_segment_words: int :param num_first_segment_words: Size of the first segment to allocate (in words ie. 8 byte increments) - :type allocate_seg_callable: Callable[[int], bytearray] + :type allocate_seg_callable: Callable[[int], Buffer] :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte - words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory - allocation strategy. + words to allocate (as an `int`) and returns any object supporting the writable buffer protocol + (e.g., `bytearray`, `memoryview`, `numpy.ndarray`). This enables custom memory allocation + strategies including shared memory. :rtype: :class:`_DynamicStructBuilder` """ @@ -1700,10 +1701,11 @@ cdef class _DynamicStructBuilder: :type num_first_segment_words: int :param num_first_segment_words: Size of the first segment to allocate (in words ie. 8 byte increments) - :type allocate_seg_callable: Callable[[int], bytearray] + :type allocate_seg_callable: Callable[[int], Buffer] :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte - words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory - allocation strategy. + words to allocate (as an `int`) and returns any object supporting the writable buffer protocol + (e.g., `bytearray`, `memoryview`, `numpy.ndarray`). This enables custom memory allocation + strategies including shared memory. :rtype: :class:`_DynamicStructBuilder` """ @@ -3891,15 +3893,27 @@ cdef class _PyCustomMessageBuilder(_MessageBuilder): This callable object will be invoked in the allocateSegment method of the MessageBuilder to allocate memory. The allocated memory will be managed within the MessageBuilder. - :type allocate_seg_callable: Callable[[int], bytearray] + :type allocate_seg_callable: Callable[[int], Buffer] :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte - words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory - allocation strategy. + words to allocate (as an `int`) and returns any object supporting the writable buffer protocol + (e.g., `bytearray`, `memoryview`, `numpy.ndarray`). This enables custom memory allocation + strategies including shared memory. Required function signature is like this: - def __call__(self, minimum_size: int) -> bytearray: + def __call__(self, minimum_size: int) -> Buffer: + + Where `Buffer` is any object that: + - Supports the Python buffer protocol (PyObject_GetBuffer) + - Is writable Note that the unit of minimum_size is words, ie. 8 byte increments. + The underlying memory must remain valid for the lifetime of the MessageBuilder. + If returning a view (e.g., `memoryview`, `numpy.ndarray`) that wraps external memory, + the allocator is responsible for properly managing the memory lifecycle。 + + Examples: + + # Example 1: Simple bytearray allocator class Allocator: def __init__(self): self.cur_size = 0 @@ -3911,9 +3925,32 @@ cdef class _PyCustomMessageBuilder(_MessageBuilder): return bytearray(byte_count) addressbook = capnp.load('addressbook.capnp') + allocator = Allocator() message = capnp._PyCustomMessageBuilder(allocator) person = message.init_root(addressbook.Person) + # Example 2: Shared memory allocator (zero-copy) + import ctypes + + class ShmAllocator: + def __init__(self, shm_pool): + self.shm = shm_pool + self.buffers = [] + + def __call__(self, minimum_size: int) -> memoryview: + size = minimum_size * 8 + ptr = self.shm.allocate(size) + buffer = (ctypes.c_uint8 * size).from_address(ptr) + self.buffers.append(buffer) + return memoryview(buffer) + + def release(self): + for buffer in self.buffers: + ptr = ctypes.addressof(buffer) + size = ctypes.sizeof(buffer) + self.shm.deallocate(ptr, size) + self.buffers.clear() + :type size: int :param size: Size of the first segment to allocate (in words ie. 8 byte increments) """ diff --git a/examples/py_custom_message_builder.py b/examples/py_custom_message_builder.py index f5e06a481..95d7aa35c 100644 --- a/examples/py_custom_message_builder.py +++ b/examples/py_custom_message_builder.py @@ -23,6 +23,28 @@ def __call__(self, minimum_size: int) -> bytearray: return bytearray(byte_count) +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) + print( + f"minimum_size: {minimum_size}, last_size: {self.last_size}, " + f"actual_size: {actual_size}, cur_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) + + person = addressbook_capnp.Person.new_message(allocate_seg_callable=Allocator()) person.init("extraData", 5) @@ -46,3 +68,28 @@ def __call__(self, minimum_size: int) -> bytearray: print(person.extraData) print(bytes(person.extraData)) print(type(person.extraData)) +print() + +print("=" * 40) +print("MemoryViewAllocator") +print("=" * 40) + +person = addressbook_capnp.Person.new_message( + allocate_seg_callable=MemoryViewAllocator() +) + +person.init("extraData", 5) +print(person.extraData) +print(bytes(person.extraData)) +print(type(person.extraData)) +print() + +person.extraData[1] = 0xFF +print(person.extraData) +print(bytes(person.extraData)) +print() + +person.extraData = b"world" +print(person.extraData) +print(bytes(person.extraData)) +print(type(person.extraData)) diff --git a/test/test_py_custom_message_builder.py b/test/test_py_custom_message_builder.py index 1b5b8f127..5011b5c2f 100644 --- a/test/test_py_custom_message_builder.py +++ b/test/test_py_custom_message_builder.py @@ -12,7 +12,7 @@ def all_types(): return capnp.load(os.path.join(this_dir, "all_types.capnp")) -def test_addressbook(all_types): +def test_bytearray_allocator(all_types): class Allocator: def __init__(self): self.cur_size = 0 @@ -20,10 +20,6 @@ def __init__(self): def __call__(self, minimum_size: int) -> bytearray: actual_size = max(minimum_size, self.cur_size) - print( - f"minimum_size: {minimum_size}, last_size: {self.last_size}, " - f"actual_size: {actual_size}, cur_size: {self.cur_size}" - ) self.last_size = actual_size self.cur_size += actual_size WORD_SIZE = 8 @@ -39,10 +35,45 @@ def __call__(self, minimum_size: int) -> bytearray: assert allocator.last_size == 1024 struct_builder.init("dataField", 5) - assert struct_builder._get("dataField") == b"\x00\x00\x00\x00\x00" + assert bytes(struct_builder._get("dataField")) == b"\x00\x00\x00\x00\x00" struct_builder.dataField = b"hello" - assert struct_builder._get("dataField") == b"hello" + assert bytes(struct_builder._get("dataField")) == b"hello" - struct_builder = struct_builder.as_reader() - assert 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" From e5802f8156d5483b1476bf291d414dccbe6d0b48 Mon Sep 17 00:00:00 2001 From: Brian Xu Date: Fri, 30 Jan 2026 00:15:30 +1100 Subject: [PATCH 2/3] refine example --- examples/py_custom_message_builder.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/examples/py_custom_message_builder.py b/examples/py_custom_message_builder.py index 95d7aa35c..0c99fa88e 100644 --- a/examples/py_custom_message_builder.py +++ b/examples/py_custom_message_builder.py @@ -49,47 +49,28 @@ def __call__(self, minimum_size: int) -> memoryview: person.init("extraData", 5) print(person.extraData) -print(bytes(person.extraData)) print(type(person.extraData)) print() -person.extraData[1] = 0xFF -print(person.extraData) -print(bytes(person.extraData)) -print() - person.extraData = b"hello" print(person.extraData) -print(bytes(person.extraData)) print(type(person.extraData)) print() person = person.as_reader() print(person.extraData) -print(bytes(person.extraData)) print(type(person.extraData)) print() -print("=" * 40) -print("MemoryViewAllocator") -print("=" * 40) - person = addressbook_capnp.Person.new_message( allocate_seg_callable=MemoryViewAllocator() ) person.init("extraData", 5) print(person.extraData) -print(bytes(person.extraData)) print(type(person.extraData)) print() -person.extraData[1] = 0xFF -print(person.extraData) -print(bytes(person.extraData)) -print() - person.extraData = b"world" print(person.extraData) -print(bytes(person.extraData)) print(type(person.extraData)) From c67243b285fc2139ecb89c3223d5240edf5c59d8 Mon Sep 17 00:00:00 2001 From: Brian Xu Date: Fri, 30 Jan 2026 00:21:48 +1100 Subject: [PATCH 3/3] refine example --- examples/py_custom_message_builder.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/examples/py_custom_message_builder.py b/examples/py_custom_message_builder.py index 0c99fa88e..c10b79b0b 100644 --- a/examples/py_custom_message_builder.py +++ b/examples/py_custom_message_builder.py @@ -25,22 +25,11 @@ def __call__(self, minimum_size: int) -> bytearray: 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) - print( - f"minimum_size: {minimum_size}, last_size: {self.last_size}, " - f"actual_size: {actual_size}, cur_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) + buffer = bytearray(minimum_size * WORD_SIZE) self.buffers.append(buffer) return memoryview(buffer)