IQ.Lvbs/pycapnp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Brian Xu
2025-09-10 23:48:58 +10:00
committed by GitHub
parent 62682977cb
commit 84674909a2
10 changed files with 304 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
capnp/lib/capnp.pxd
View File

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

130
capnp/lib/capnp.pyx
View File

@@ -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.
@@ -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(<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