diff --git a/capnp/lib/capnp.pxd b/capnp/lib/capnp.pxd index ed5961fae..a4b1da7d7 100644 --- a/capnp/lib/capnp.pxd +++ b/capnp/lib/capnp.pxd @@ -82,6 +82,7 @@ cdef class _DynamicStructBuilder: cdef _check_write(self) cpdef to_bytes(_DynamicStructBuilder self) cpdef to_segments(_DynamicStructBuilder self) + cpdef to_segment_views(_DynamicStructBuilder self) cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) cpdef to_bytes_packed(_DynamicStructBuilder self) diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index a6eddb0f3..7230d6a01 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -19,6 +19,7 @@ from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_REA from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer from cpython.bytes cimport PyBytes_FromStringAndSize from cpython.exc cimport PyErr_Clear +from cpython.pyport cimport PY_SSIZE_T_MAX from cython.operator cimport dereference as deref from libc.stdlib cimport malloc, free from libc.string cimport memcpy @@ -1173,6 +1174,70 @@ cdef class _MessageSize: self.word_count = word_count self.cap_count = cap_count + +@cython.internal +cdef class _SegmentView: + cdef object _builder + cdef const char* _ptr + cdef Py_ssize_t _size + + cdef _init(self, object builder, const char* ptr, Py_ssize_t size): + self._builder = builder + self._ptr = ptr + self._size = size + return self + + def __getbuffer__(self, Py_buffer *buffer, int flags): + if PyBuffer_FillInfo(buffer, self, self._ptr, self._size, 1, flags) < 0: + raise BufferError("Failed to create segment buffer view") + + def __releasebuffer__(self, Py_buffer *buffer): + pass + + def __len__(self): + return self._size + + def __repr__(self): + return '' % self._size + + +@cython.internal +cdef class _SegmentViews: + cdef object _builder + cdef list _views + + cdef _init(self, _MessageBuilder builder): + cdef schema_cpp.ConstWordArrayArrayPtr segments = builder.thisptr.getSegmentsForOutput() + cdef size_t i + cdef size_t word_count + cdef Py_ssize_t byte_count + + self._builder = builder + self._views = [] + for i in range(0, segments.size()): + word_count = segments[i].size() + if word_count > (PY_SSIZE_T_MAX // 8): + raise OverflowError("segment is too large to expose as a Python buffer") + byte_count = (8 * word_count) + self._views.append(_SegmentView()._init( + builder, + segments[i].begin(), + byte_count)) + return self + + def __getitem__(self, index): + return self._views[index] + + def __iter__(self): + return iter(self._views) + + def __len__(self): + return len(self._views) + + def __repr__(self): + return '' % len(self) + + if getattr(_sys, 'subversion', [''])[0] == 'PyPy': from pickle_helper import _struct_reducer else: @@ -1451,7 +1516,8 @@ cdef class _DynamicStructBuilder: cpdef to_segments(_DynamicStructBuilder self): """Returns the struct's containing message as a Python list of Python bytes objects. - This avoids making copies. + This copies each output segment into a Python-owned bytes object. Use + to_segment_views() for zero-copy, read-only borrowed segment views. NB: This is not currently supported on PyPy. @@ -1462,6 +1528,18 @@ cdef class _DynamicStructBuilder: segments = builder.get_segments_for_output() return segments + cpdef to_segment_views(_DynamicStructBuilder self): + """Returns the struct's containing message as zero-copy, read-only segment views. + + The returned views borrow memory from the message builder. Do not mutate, reset, or reuse + the builder while the views are still in use. + + :rtype: sequence + """ + self._check_write() + cdef _MessageBuilder builder = self._parent + return _SegmentViews()._init(builder) + cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count): cdef _MessageBuilder builder = self._parent array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count) diff --git a/docs/quickstart.rst b/docs/quickstart.rst index d2d08b03a..c2e3114ec 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -338,10 +338,18 @@ For compatibility on the Python side, use the ``to_segments()`` and ``from_segme segments = alice.to_segments() -This returns a list of segments, each a byte buffer. Each segment can be, e.g., turned into a ZeroMQ message frame. The list of segments can also be turned back into an object:: +This returns a list of copied, Python-owned ``bytes`` objects. Each segment can be, e.g., turned into a ZeroMQ message frame. The list of segments can also be turned back into an object:: alice = addressbook_capnp.Person.from_segments(segments) +For high-throughput code that can safely consume borrowed buffers, ``to_segment_views()`` exposes the same output segments without copying them into Python ``bytes`` objects:: + + segment_views = alice.to_segment_views() + for segment in segment_views: + transport.send(segment) + +Each segment view supports the Python buffer protocol and is read-only. The returned views borrow memory from the message builder's arena, so do not mutate, reset, or reuse the builder while any segment view is still in use. If you need data that remains independent of the builder lifetime, use ``to_segments()`` instead. + For more information, please refer to the following links: - `Advice on minimizing copies from Cap'n Proto `_ (from the author of Cap'n Proto) diff --git a/test/test_serialization.py b/test/test_serialization.py index 6c4bfa5cc..ecf72da1c 100644 --- a/test/test_serialization.py +++ b/test/test_serialization.py @@ -1,6 +1,7 @@ import warnings from contextlib import contextmanager +import gc import pytest import capnp import os @@ -62,6 +63,93 @@ def test_roundtrip_segments(all_types): test_regression.check_all_types(msg) +@pytest.mark.skipif( + platform.python_implementation() == "PyPy", + reason="TODO: Investigate segmented serialization support on PyPy.", +) +def test_segment_views_are_read_only_buffers(all_types): + msg = all_types.TestAllTypes.new_message() + test_regression.init_all_types(msg) + + segments = msg.to_segments() + segment_views = msg.to_segment_views() + + assert len(segment_views) == len(segments) + assert len(segment_views) >= 1 + + for segment_view, segment_bytes in zip(segment_views, segments): + assert not isinstance(segment_view, bytes) + view = memoryview(segment_view) + try: + assert view.readonly is True + assert view.tobytes() == segment_bytes + finally: + view.release() + + +@pytest.mark.skipif( + platform.python_implementation() == "PyPy", + reason="TODO: Investigate segmented serialization support on PyPy.", +) +def test_roundtrip_segment_views(all_types): + msg = all_types.TestAllTypes.new_message() + test_regression.init_all_types(msg) + + segment_views = msg.to_segment_views() + msg = all_types.TestAllTypes.from_segments(segment_views) + test_regression.check_all_types(msg) + + +@pytest.mark.skipif( + platform.python_implementation() == "PyPy", + reason="TODO: Investigate segmented serialization support on PyPy.", +) +def test_segment_views_are_not_writable(all_types): + msg = all_types.TestAllTypes.new_message() + test_regression.init_all_types(msg) + + segment_views = msg.to_segment_views() + view = memoryview(segment_views[0]) + try: + assert len(view) > 0 + with pytest.raises(TypeError): + view[0] = 0 + finally: + view.release() + + +@pytest.mark.skipif( + platform.python_implementation() == "PyPy", + reason="TODO: Investigate segmented serialization support on PyPy.", +) +def test_segment_view_keeps_message_alive(all_types): + msg = all_types.TestAllTypes.new_message() + test_regression.init_all_types(msg) + + segment_views = msg.to_segment_views() + segment_view = segment_views[0] + view = memoryview(segment_view) + expected = view.tobytes() + + del msg + del segment_views + del segment_view + gc.collect() + + try: + assert view.tobytes() == expected + finally: + view.release() + + +def test_segment_views_require_root_struct(all_types): + msg = all_types.TestAllTypes.new_message() + nested = msg.init("structField") + + with pytest.raises(capnp.KjException): + nested.to_segment_views() + + @pytest.mark.skipif( sys.version_info[0] < 3, reason="mmap doesn't implement the buffer interface under python 2.",