diff --git a/capnp/includes/schema_cpp.pxd b/capnp/includes/schema_cpp.pxd index 3716648..d9b691c 100644 --- a/capnp/includes/schema_cpp.pxd +++ b/capnp/includes/schema_cpp.pxd @@ -667,6 +667,8 @@ cdef extern from "capnp/message.h" namespace " ::capnp": DynamicStruct_Builder initRootDynamicStruct'initRoot< ::capnp::DynamicStruct>'(StructSchema) void setRootDynamicStruct'setRoot< ::capnp::DynamicStruct::Reader>'(DynamicStruct.Reader) + ConstWordArrayArrayPtr getSegmentsForOutput'getSegmentsForOutput'() + AnyPointer.Builder getRootAnyPointer'getRoot< ::capnp::AnyPointer>'() DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema) @@ -691,6 +693,10 @@ cdef extern from "capnp/message.h" namespace " ::capnp": MallocMessageBuilder() MallocMessageBuilder(int) + cdef cppclass SegmentArrayMessageReader(MessageReader): + SegmentArrayMessageReader(ConstWordArrayArrayPtr array) except +reraise_kj_exception + SegmentArrayMessageReader(ConstWordArrayArrayPtr array, ReaderOptions) except +reraise_kj_exception + cdef cppclass FlatMessageBuilder(MessageBuilder): FlatMessageBuilder(WordArrayPtr array) FlatMessageBuilder(WordArrayPtr array, ReaderOptions) @@ -714,6 +720,16 @@ cdef extern from "kj/common.h" namespace " ::kj": ByteArrayPtr(byte *, size_t size) size_t size() byte& operator[](size_t index) + cdef cppclass ConstWordArrayPtr " ::kj::ArrayPtr< const ::capnp::word>": + ConstWordArrayPtr() + ConstWordArrayPtr(word *, size_t size) + size_t size() + const word* begin() + cdef cppclass ConstWordArrayArrayPtr " ::kj::ArrayPtr< const ::kj::ArrayPtr< const ::capnp::word>>": + ConstWordArrayArrayPtr() + ConstWordArrayArrayPtr(ConstWordArrayPtr*, size_t size) + size_t size() + ConstWordArrayPtr& operator[](size_t index) cdef extern from "kj/array.h" namespace " ::kj": # Cython can't handle Array[word] as a function argument diff --git a/capnp/lib/capnp.pxd b/capnp/lib/capnp.pxd index 1673365..6ac296a 100644 --- a/capnp/lib/capnp.pxd +++ b/capnp/lib/capnp.pxd @@ -53,6 +53,7 @@ cdef class _DynamicStructBuilder: cdef _check_write(self) cpdef to_bytes(_DynamicStructBuilder self) except +reraise_kj_exception + cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception cpdef to_bytes_packed(_DynamicStructBuilder self) except +reraise_kj_exception diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index a052146..8ee1c3d 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -1192,6 +1192,12 @@ cdef class _DynamicStructBuilder: self._is_written = True return ret + cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception: + self._check_write() + cdef _MessageBuilder builder = self._parent + segments = builder.get_segments_for_output() + return segments + cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception: cdef _MessageBuilder builder = self._parent array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count) @@ -2989,6 +2995,9 @@ class _StructModule(object): else: message = _FlatArrayMessageReader(buf, traversal_limit_in_words, nesting_limit) return message.get_root(self.schema) + def from_segments(self, segments): + message = _SegmentArrayMessageReader(segments) + return message.get_root(self.schema) def from_bytes_packed(self, buf, traversal_limit_in_words = None, nesting_limit = None): """Returns a Reader for the packed object in buf. @@ -3285,6 +3294,18 @@ cdef class _MessageBuilder: self.thisptr.setRootDynamicStruct((<_DynamicStructReader>value).thisptr) return self.get_root(value.schema) + cpdef get_segments_for_output(self) except +reraise_kj_exception: + segments = self.thisptr.getSegmentsForOutput() + res = [] + cdef const char* ptr + cdef bytes segment_bytes + for i in range(0, segments.size()): + segment = segments[i] + ptr = segment.begin() + segment_bytes = ptr[:8*segment.size()] + res.append(segment_bytes) + return res + cpdef new_orphan(self, schema) except +reraise_kj_exception: """A method for instantiating Cap'n Proto orphans @@ -3738,6 +3759,36 @@ cdef class _FlatArrayMessageReader(_MessageReader): del self.thisptr +@cython.internal +cdef class _SegmentArrayMessageReader(_MessageReader): + + cdef object _objects_to_pin + + def __init__(self, segments): + # takes a Python array of bytes and constructs a ConstWordArrayArrayPtr + num_segments = len(segments) + cdef char* ptr + cdef schema_cpp.ConstWordArrayPtr seg_ptr + cdef schema_cpp.ConstWordArrayPtr* seg_ptrs = malloc(num_segments * sizeof(schema_cpp.ConstWordArrayPtr)) + self._objects_to_pin = [] + for i in range(0, num_segments): + segment = bytes(segments[i]) + ptr = segment + if (ptr) % 8 != 0: + aligned = _AlignedBuffer(segment) + ptr = aligned.buf + self._objects_to_pin.append(aligned) + else: + self._objects_to_pin.append(segment) + seg_ptr = schema_cpp.ConstWordArrayPtr(ptr, len(segment)//8) + seg_ptrs[i] = seg_ptr + self.thisptr = new schema_cpp.SegmentArrayMessageReader( + schema_cpp.ConstWordArrayArrayPtr(seg_ptrs, num_segments)) + + def __dealloc__(self): + del self.thisptr + + @cython.internal cdef class _FlatMessageBuilder(_MessageBuilder): cdef object _object_to_pin