refine documentation for PyCustomMessageBuilder (#395)

* refine document for py custom message builder

* refine example

* refine example
This commit is contained in:
Brian Xu
2026-02-01 08:15:56 +11:00
committed by GitHub
parent aa90ab6844
commit 76a41c8efe
3 changed files with 112 additions and 27 deletions

View File

@@ -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)
"""

View File

@@ -23,26 +23,43 @@ class Allocator:
return bytearray(byte_count)
class MemoryViewAllocator:
def __init__(self):
self.buffers = []
def __call__(self, minimum_size: int) -> memoryview:
WORD_SIZE = 8
buffer = bytearray(minimum_size * WORD_SIZE)
self.buffers.append(buffer)
return memoryview(buffer)
person = addressbook_capnp.Person.new_message(allocate_seg_callable=Allocator())
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()
person = addressbook_capnp.Person.new_message(
allocate_seg_callable=MemoryViewAllocator()
)
person.init("extraData", 5)
print(person.extraData)
print(type(person.extraData))
print()
person.extraData = b"world"
print(person.extraData)
print(type(person.extraData))

View File

@@ -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 test_addressbook(all_types):
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 test_addressbook(all_types):
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"