refine documentation for PyCustomMessageBuilder (#395)

* refine document for py custom message builder

* refine example

* refine example

Author: BrianXu0623

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

examples/py_custom_message_builder.py

@@ -23,26 +23,43 @@ def __call__(self, minimum_size: int) -> bytearray:
         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
+person.extraData = b"hello"
 print(person.extraData)
-print(bytes(person.extraData))
+print(type(person.extraData))
 print()
 
-person.extraData = b"hello"
+person = person.as_reader()
 print(person.extraData)
-print(bytes(person.extraData))
 print(type(person.extraData))
 print()
 
-person = person.as_reader()
+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(bytes(person.extraData))
 print(type(person.extraData))

test/test_py_custom_message_builder.py

@@ -12,18 +12,14 @@ 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
             self.last_size = 0
 
         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"