From 84674909a29042ca3bcd938eb6230fb25acda939 Mon Sep 17 00:00:00 2001 From: Brian Xu Date: Wed, 10 Sep 2025 23:48:58 +1000 Subject: [PATCH] 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 --- capnp/__init__.py | 1 + capnp/includes/PyCustomMessageBuilder.cpp | 50 +++++++++ capnp/includes/PyCustomMessageBuilder.h | 28 +++++ capnp/includes/schema_cpp.pxd | 5 + capnp/lib/capnp.pxd | 4 +- capnp/lib/capnp.pyx | 130 +++++++++++++++++++--- examples/py_custom_message_builder.py | 48 ++++++++ setup.py | 1 + test/test_blob_to_dict_base64.py | 3 +- test/test_py_custom_message_builder.py | 51 +++++++++ 10 files changed, 304 insertions(+), 17 deletions(-) create mode 100644 capnp/includes/PyCustomMessageBuilder.cpp create mode 100644 capnp/includes/PyCustomMessageBuilder.h create mode 100644 examples/py_custom_message_builder.py create mode 100644 test/test_py_custom_message_builder.py diff --git a/capnp/__init__.py b/capnp/__init__.py index 84cfeb1..9d10c1e 100644 --- a/capnp/__init__.py +++ b/capnp/__init__.py @@ -48,6 +48,7 @@ from .lib.capnp import ( _InterfaceModule, _ListSchema, _MallocMessageBuilder, + _PyCustomMessageBuilder, _PackedFdMessageReader, _StreamFdMessageReader, _StructModule, diff --git a/capnp/includes/PyCustomMessageBuilder.cpp b/capnp/includes/PyCustomMessageBuilder.cpp new file mode 100644 index 0000000..9a08833 --- /dev/null +++ b/capnp/includes/PyCustomMessageBuilder.cpp @@ -0,0 +1,50 @@ +#include "PyCustomMessageBuilder.h" +#include + +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 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(view.buf), wordCount); +} + +} \ No newline at end of file diff --git a/capnp/includes/PyCustomMessageBuilder.h b/capnp/includes/PyCustomMessageBuilder.h new file mode 100644 index 0000000..0aba128 --- /dev/null +++ b/capnp/includes/PyCustomMessageBuilder.h @@ -0,0 +1,28 @@ +#pragma once + +#include "Python.h" +#include +#include +#include + +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 allocateSegment(capnp::uint minimumSize) override; + +private: + PyObject* allocateSegmentCallable; + + uint firstSize; + uint curSize = 0; + + std::vector allocatedBuffers; +}; + +} diff --git a/capnp/includes/schema_cpp.pxd b/capnp/includes/schema_cpp.pxd index f52add9..4d869cb 100644 --- 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 diff --git a/capnp/lib/capnp.pxd b/capnp/lib/capnp.pxd index 1c06a5c..d845767 100644 --- 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 diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 34fc63c..b69e092 100644 --- 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 (temp_text.begin())[:temp_text.size()] elif type == capnp.TYPE_DATA: temp_data = self.asData() - return ((temp_data.begin())[:temp_data.size()]) + return PyMemoryView_FromMemory( 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 (temp_text.begin())[:temp_text.size()] elif type == capnp.TYPE_DATA: temp_data = self.asData() - return ((temp_data.begin())[:temp_data.size()]) + return PyMemoryView_FromMemory( 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( 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(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. @@ -1250,11 +1283,20 @@ 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` """ - builder = _MallocMessageBuilder(num_first_segment_words) - return builder.set_root(self) + 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. @@ -1601,11 +1643,20 @@ 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` """ - builder = _MallocMessageBuilder(num_first_segment_words) - return builder.set_root(self) + 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): - builder = _MallocMessageBuilder(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(allocate_seg_callable) + else: + self.thisptr = new schema_cpp.PyCustomMessageBuilder(allocate_seg_callable, size) + + cdef class _MessageReader: """An abstract base class for reading Cap'n Proto messages diff --git a/examples/py_custom_message_builder.py b/examples/py_custom_message_builder.py new file mode 100644 index 0000000..f5e06a4 --- /dev/null +++ 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)) diff --git a/setup.py b/setup.py index 46864b0..e5c342c 100644 --- 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, diff --git a/test/test_blob_to_dict_base64.py b/test/test_blob_to_dict_base64.py index 5e70f7e..c6f7a03 100644 --- 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 diff --git a/test/test_py_custom_message_builder.py b/test/test_py_custom_message_builder.py new file mode 100644 index 0000000..1b62e1a --- /dev/null +++ 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"