Support python custom message builder and make Data field's type return MemoryView (#380)

This PR is for resolving the following issue: [issue](https://github.com/capnproto/pycapnp/issues/379) 1. Created `_PyCustomMessageBuilder` extends `MessageBuilder`, enabling the ability to customise the `SegmentAllocate` method in Python. This allows allocation and data population within shared memory, and supports zero-copy inter-process data transfer by passing segment offsets. 2. Fields of type `Data` now support being set with a `memoryview`. When retrieving a `Data` field from a `DynamicStructBuilder`, it will return a writable `memoryview`, allowing users to modify the data directly. This enables memory to be pre-allocated and content to be modified in later, eliminating an extra copy. When retrieving a `Data` field from a `DynamicStructReader`, it will return a read-only `memoryview`, allowing user to read data without memory copy. * add memoryview and custom builder * support set dynamic field * add curSize * add initialSize and lastSize * change StringPtr name * add test case * refine test case * convert func to py callable object * add initial value * refine example * add copy as_reader and new_message, make structReader's data field return RO memoryView * rebase master and bugfix * reformat flake8 * refine test case * refine test cases for blob * remove unused import for flake8 * run black . --------- Co-authored-by: Brian Xu <brian.xu1@bytedance.com>
2025-09-10 23:48:58 +10:00
parent 62682977cb
commit 84674909a2
10 changed files with 304 additions and 17 deletions
--- a/capnp/init.py
+++ b/capnp/init.py
@@ -48,6 +48,7 @@ from .lib.capnp import (
    _InterfaceModule,
    _ListSchema,
    _MallocMessageBuilder,
+    _PyCustomMessageBuilder,
    _PackedFdMessageReader,
    _StreamFdMessageReader,
    _StructModule,
--- a/capnp/includes/PyCustomMessageBuilder.cpp
+++ b/capnp/includes/PyCustomMessageBuilder.cpp
@@ -0,0 +1,50 @@
+#include "PyCustomMessageBuilder.h"
+#include <stdexcept>
+
+namespace capnp {
+
+PyCustomMessageBuilder::PyCustomMessageBuilder(
+    PyObject* allocateSegmentCallable, uint firstSegmentWords)
+    : allocateSegmentCallable(allocateSegmentCallable), firstSize(firstSegmentWords)
+{
+  KJ_REQUIRE(PyCallable_Check(allocateSegmentCallable),
+               "allocateSegmentCallable must be callable");
+  Py_INCREF(allocateSegmentCallable);
+}
+
+PyCustomMessageBuilder::~PyCustomMessageBuilder() noexcept(false) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+
+  for (auto* obj : allocatedBuffers) {
+    Py_DECREF(obj);
+  }
+  allocatedBuffers.clear();
+
+  Py_DECREF(allocateSegmentCallable);
+  PyGILState_Release(gstate);
+}
+
+kj::ArrayPtr<capnp::word> PyCustomMessageBuilder::allocateSegment(capnp::uint minimumSize) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+  KJ_DEFER({ PyGILState_Release(gstate); });
+  if (curSize == 0) {
+    minimumSize = kj::max(minimumSize, firstSize);
+  }
+  PyObject* pyBufObj = PyObject_CallFunction(allocateSegmentCallable, "I", minimumSize);
+  KJ_REQUIRE(pyBufObj, "PyCustomMessageBuilder: allocateSegment failed");
+  allocatedBuffers.push_back(pyBufObj);
+
+
+  Py_buffer view;
+  int bufRes = PyObject_GetBuffer(pyBufObj, &view, PyBUF_SIMPLE);
+  KJ_REQUIRE(bufRes == 0, "PyCustomMessageBuilder: object does not support buffer protocol");
+  KJ_DEFER({ PyBuffer_Release(&view); });
+
+  size_t byteCount = view.len;
+  size_t wordCount = byteCount / sizeof(capnp::word);
+  KJ_REQUIRE(wordCount >= minimumSize, "PyCustomMessageBuilder: buffer too small for minimumSize");
+  curSize += wordCount;
+  return kj::arrayPtr(reinterpret_cast<capnp::word*>(view.buf), wordCount);
+}
+
+}
--- a/capnp/includes/PyCustomMessageBuilder.h
+++ b/capnp/includes/PyCustomMessageBuilder.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include "Python.h"
+#include <capnp/message.h>
+#include <capnp/serialize.h>
+#include <vector>
+
+namespace capnp {
+
+class PyCustomMessageBuilder : public capnp::MessageBuilder {
+public:
+  explicit PyCustomMessageBuilder(PyObject* allocateSegmentCallable,
+  uint firstSegmentWords = capnp::SUGGESTED_FIRST_SEGMENT_WORDS);
+
+  ~PyCustomMessageBuilder() noexcept(false) override;
+
+  kj::ArrayPtr<capnp::word> allocateSegment(capnp::uint minimumSize) override;
+
+private:
+  PyObject* allocateSegmentCallable;
+
+  uint firstSize;
+  uint curSize = 0;
+
+  std::vector<PyObject*> allocatedBuffers;
+};
+
+}
--- a/capnp/includes/schema_cpp.pxd
+++ b/capnp/includes/schema_cpp.pxd
@@ -714,6 +714,11 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
    enum Void:
        VOID

+cdef extern from "PyCustomMessageBuilder.h" namespace " ::capnp":
+    cdef cppclass PyCustomMessageBuilder(MessageBuilder):
+        PyCustomMessageBuilder(PyObject* allocateSegmentCallable)
+        PyCustomMessageBuilder(PyObject* allocateSegmentCallable, int firstSegmentSize)
+
 cdef extern from "capnp/common.h" namespace " ::capnp":
    cdef cppclass word nogil:
        pass
--- a/capnp/lib/capnp.pxd
+++ b/capnp/lib/capnp.pxd
@@ -66,7 +66,7 @@ cdef class _DynamicStructReader:
    cpdef _get_by_field(self, _StructSchemaField field)
    cpdef _has_by_field(self, _StructSchemaField field)

-    cpdef as_builder(self, num_first_segment_words=?)
+    cpdef as_builder(self, num_first_segment_words=?, allocate_seg_callable=?)


 cdef class _DynamicStructBuilder:
@@ -99,7 +99,7 @@ cdef class _DynamicStructBuilder:
    cpdef disown(self, field)

    cpdef as_reader(self)
-    cpdef copy(self, num_first_segment_words=?)
+    cpdef copy(self, num_first_segment_words=?, allocate_seg_callable=?)

 cdef class _DynamicEnumField:
    cdef object thisptr
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -13,8 +13,9 @@ from capnp.helpers.helpers cimport init_capnp_api
 from capnp.includes.capnp_cpp cimport AsyncIoStream, WaitScope, PyPromise, VoidPromise, EventPort, EventLoop, PyAsyncIoStream, PromiseFulfiller, VoidPromiseFulfiller, tryReadMessage, writeMessage, makeException, PythonInterfaceDynamicImpl
 from capnp.includes.schema_cpp cimport (MessageReader,)

+from builtins import memoryview as BuiltinsMemoryview
 from cpython cimport array, Py_buffer, PyObject_CheckBuffer
-from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ
+from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO
 from cpython.memoryview cimport PyMemoryView_FromMemory
 from cpython.exc cimport PyErr_Clear
 from cython.operator cimport dereference as deref
@@ -667,7 +668,7 @@ cdef to_python_reader(C_DynamicValue.Reader self, object parent):
        return (<char*>temp_text.begin())[:temp_text.size()]
    elif type == capnp.TYPE_DATA:
        temp_data = self.asData()
-        return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
+        return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_READ)
    elif type == capnp.TYPE_LIST:
        return _DynamicListReader()._init(self.asList(), parent)
    elif type == capnp.TYPE_STRUCT:
@@ -701,7 +702,7 @@ cdef to_python_builder(C_DynamicValue.Builder self, object parent):
        return (<char*>temp_text.begin())[:temp_text.size()]
    elif type == capnp.TYPE_DATA:
        temp_data = self.asData()
-        return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
+        return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_WRITE)
    elif type == capnp.TYPE_LIST:
        return _DynamicListBuilder()._init(self.asList(), parent)
    elif type == capnp.TYPE_STRUCT:
@@ -766,6 +767,20 @@ cdef _setBytes(_DynamicSetterClasses thisptr, field, value):
    cdef C_DynamicValue.Reader temp = C_DynamicValue.Reader(temp_string)
    thisptr.set(field, temp)

+cdef _setMemoryview(_DynamicSetterClasses thisptr, field, value):
+    cdef Py_buffer buf
+    cdef capnp.StringPtr temp_string
+    cdef C_DynamicValue.Reader temp
+    if PyObject_GetBuffer(value, &buf, PyBUF_CONTIG_RO) != 0:
+        raise KjException(
+            "cannot get buffer from memory view, for field '{}'".format(field)
+        )
+    try:
+        temp_string = capnp.StringPtr(<char *> buf.buf, buf.len)
+        temp = C_DynamicValue.Reader(temp_string)
+        thisptr.set(field, temp)
+    finally:
+        PyBuffer_Release(&buf)

 cdef _setBaseString(_DynamicSetterClasses thisptr, field, value):
    encoded_value = value.encode('utf-8')
@@ -779,6 +794,20 @@ cdef _setBytesField(DynamicStruct_Builder thisptr, _StructSchemaField field, val
    cdef C_DynamicValue.Reader temp = C_DynamicValue.Reader(temp_string)
    thisptr.setByField(field.thisptr, temp)

+cdef _setMemoryviewField(DynamicStruct_Builder thisptr, _StructSchemaField field, value):
+    cdef Py_buffer buf
+    cdef capnp.StringPtr temp_string
+    cdef C_DynamicValue.Reader temp
+    if PyObject_GetBuffer(value, &buf, PyBUF_CONTIG_RO) != 0:
+        raise KjException(
+            "cannot get buffer from memory view, for field '{}'".format(field)
+        )
+    try:
+        temp_string = capnp.StringPtr(<char *>buf.buf, buf.len)
+        temp = C_DynamicValue.Reader(temp_string)
+        thisptr.setByField(field.thisptr, temp)
+    finally:
+        PyBuffer_Release(&buf)

 cdef _setBaseStringField(DynamicStruct_Builder thisptr, _StructSchemaField field, value):
    encoded_value = value.encode('utf-8')
@@ -805,6 +834,8 @@ cdef _setDynamicField(_DynamicSetterClasses thisptr, field, value, parent):
        thisptr.set(field, temp)
    elif value_type is bytes:
        _setBytes(thisptr, field, value)
+    elif isinstance(value, BuiltinsMemoryview):
+        _setMemoryview(thisptr, field, value)
    elif isinstance(value, basestring):
        _setBaseString(thisptr, field, value)
    elif value_type is list:
@@ -870,6 +901,8 @@ cdef _setDynamicFieldWithField(DynamicStruct_Builder thisptr, _StructSchemaField
        thisptr.setByField(field.thisptr, temp)
    elif value_type is bytes:
        _setBytesField(thisptr, field, value)
+    elif isinstance(value, BuiltinsMemoryview):
+        _setMemoryviewField(thisptr, field, value)
    elif isinstance(value, basestring):
        _setBaseStringField(thisptr, field, value)
    elif value_type is list:
@@ -1242,7 +1275,7 @@ cdef class _DynamicStructReader:
    def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
        return _to_dict(self, verbose, ordered, encode_bytes_as_base64)

-    cpdef as_builder(self, num_first_segment_words=None):
+    cpdef as_builder(self, num_first_segment_words=None, allocate_seg_callable=None):
        """A method for casting this Reader to a Builder

        This is a copying operation with respect to the message's buffer.
@@ -1251,10 +1284,19 @@ 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]
+        :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.
+
        :rtype: :class:`_DynamicStructBuilder`
        """
+        if allocate_seg_callable is None:
            builder = _MallocMessageBuilder(num_first_segment_words)
            return builder.set_root(self)
+        else:
+            builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
+            return builder.set_root(self)

    property total_size:
        def __get__(self):
@@ -1593,7 +1635,7 @@ cdef class _DynamicStructBuilder:
        reader._obj_to_pin = self
        return reader

-    cpdef copy(self, num_first_segment_words=None):
+    cpdef copy(self, num_first_segment_words=None, allocate_seg_callable=None):
        """A method for copying this Builder

        This is a copying operation with respect to the message's buffer.
@@ -1602,10 +1644,19 @@ 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]
+        :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.
+
        :rtype: :class:`_DynamicStructBuilder`
        """
+        if allocate_seg_callable is None:
            builder = _MallocMessageBuilder(num_first_segment_words)
            return builder.set_root(self)
+        else:
+            builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
+            return builder.set_root(self)

    property schema:
        """A property that returns the _StructSchema object matching this writer"""
@@ -3145,8 +3196,12 @@ class _StructABCMeta(type):
        return isinstance(obj, cls.__base__) and obj.schema == cls._schema


-cdef _new_message(self, kwargs, num_first_segment_words):
+cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable):
+    cdef _MessageBuilder builder
+    if allocate_seg_callable is None:
        builder = _MallocMessageBuilder(num_first_segment_words)
+    else:
+        builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
    msg = builder.init_root(self.schema)
    if kwargs is not None:
        msg.from_dict(kwargs)
@@ -3387,12 +3442,17 @@ class _StructModule(object):
    def __call__(self, num_first_segment_words=None, **kwargs):
        return self.new_message(num_first_segment_words=num_first_segment_words, **kwargs)

-    def new_message(self, num_first_segment_words=None, **kwargs):
+    def new_message(self, num_first_segment_words=None, allocate_seg_callable=None, **kwargs):
        """Returns a newly allocated builder message.

        :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]
+        :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.
+
        :type kwargs: dict
        :param kwargs: A list of fields and their values to initialize in the struct.

@@ -3401,7 +3461,7 @@ class _StructModule(object):

        :rtype: :class:`_DynamicStructBuilder`
        """
-        return _new_message(self, kwargs, num_first_segment_words)
+        return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable)


 class _InterfaceModule(object):
@@ -3758,6 +3818,50 @@ cdef class _MallocMessageBuilder(_MessageBuilder):
            self.thisptr = new schema_cpp.MallocMessageBuilder(size)


+cdef class _PyCustomMessageBuilder(_MessageBuilder):
+    """The class for building Cap'n Proto messages,
+    with customised memory allocation strategy
+
+    You will use this class if you want to customise the allocateSegment method,
+    and define your own memory allocation strategy.
+    """
+    def __init__(self, allocate_seg_callable, size=None):
+        """ The constructor requires you to provide a Python callable object as a parameter.
+        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]
+        :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.
+
+        Required function signature is like this:
+        def __call__(self, minimum_size: int) -> bytearray:
+        Note that the unit of minimum_size is words, ie. 8 byte increments.
+
+            class Allocator:
+                def __init__(self):
+                    self.cur_size = 0
+                def __call__(self, minimum_size: int) -> bytearray:
+                    size = max(minimum_size, self.cur_size)
+                    self.cur_size += size
+                    WORD_SIZE = 8
+                    byte_count = size * WORD_SIZE
+                    return bytearray(byte_count)
+
+            addressbook = capnp.load('addressbook.capnp')
+            message = capnp._PyCustomMessageBuilder(allocator)
+            person = message.init_root(addressbook.Person)
+
+        :type size: int
+        :param size: Size of the first segment to allocate (in words ie. 8 byte increments)
+        """
+        if size is None:
+            self.thisptr = new schema_cpp.PyCustomMessageBuilder(<PyObject*>allocate_seg_callable)
+        else:
+            self.thisptr = new schema_cpp.PyCustomMessageBuilder(<PyObject*>allocate_seg_callable, size)
+
+
 cdef class _MessageReader:
    """An abstract base class for reading Cap'n Proto messages

--- a/examples/py_custom_message_builder.py
+++ b/examples/py_custom_message_builder.py
@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+import capnp  # noqa: F401
+import addressbook_capnp
+
+
+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
+        byte_count = actual_size * WORD_SIZE
+        return bytearray(byte_count)
+
+
+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))
--- a/setup.py
+++ b/setup.py
@@ -201,6 +201,7 @@ extensions = [
        "*",
        [
            "capnp/helpers/capabilityHelper.cpp",
+            "capnp/includes/PyCustomMessageBuilder.cpp",
            "capnp/lib/*.pyx",
        ],
        extra_compile_args=extra_compile_args,
--- a/test/test_blob_to_dict_base64.py
+++ b/test/test_blob_to_dict_base64.py
@@ -1,6 +1,5 @@
 import os
 import capnp
-import base64
 import pytest

 this_dir = os.path.dirname(__file__)
@@ -15,7 +14,7 @@ def test_blob_to_dict(blob_schema):
    blob_value = b"hello world"
    blob = blob_schema.BlobTest(blob=blob_value)
    blob_dict = blob.to_dict(encode_bytes_as_base64=True)
-    assert base64.b64decode(blob_dict["blob"]) == blob_value
+    assert blob_dict["blob"].tobytes() == blob_value
    msg = blob_schema.BlobTest.new_message()
    msg.from_dict(blob_dict)
    assert blob.blob == blob_value
--- a/test/test_py_custom_message_builder.py
+++ b/test/test_py_custom_message_builder.py
@@ -0,0 +1,51 @@
+#!/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_addressbook(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
+            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 struct_builder._get("dataField") == b"\x00\x00\x00\x00\x00"
+
+    struct_builder._get("dataField")[1] = 0xFF
+    assert struct_builder._get("dataField") == b"\x00\xff\x00\x00\x00"
+
+    struct_builder.dataField = b"hello"
+    assert struct_builder._get("dataField") == b"hello"
+
+    struct_builder = struct_builder.as_reader()
+    assert struct_builder._get("dataField") == b"hello"