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>
This commit is contained in:
50
capnp/includes/PyCustomMessageBuilder.cpp
Normal file
50
capnp/includes/PyCustomMessageBuilder.cpp
Normal file
@@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user