From 9237ba123a043e253d271d680ac63be9afdc1839 Mon Sep 17 00:00:00 2001 From: Brian Xu Date: Fri, 16 Jan 2026 16:40:22 +1100 Subject: [PATCH] Revert Data fields to bytes and add get_data_as_view for zero-copy access (#390) * get data field with view * refine tc * refine based on flake check * run black again * rebase upstream master * add comment to tc * refine raise exception --- capnp/lib/capnp.pxd | 2 + capnp/lib/capnp.pyx | 57 ++++++- test/test_get_data_view.py | 210 +++++++++++++++++++++++++ test/test_py_custom_message_builder.py | 3 - 4 files changed, 265 insertions(+), 7 deletions(-) create mode 100644 test/test_get_data_view.py diff --git a/capnp/lib/capnp.pxd b/capnp/lib/capnp.pxd index d845767..ed5961f 100644 --- a/capnp/lib/capnp.pxd +++ b/capnp/lib/capnp.pxd @@ -65,6 +65,7 @@ cdef class _DynamicStructReader: cpdef _which_str(self) cpdef _get_by_field(self, _StructSchemaField field) cpdef _has_by_field(self, _StructSchemaField field) + cpdef get_data_as_view(self, field) cpdef as_builder(self, num_first_segment_words=?, allocate_seg_callable=?) @@ -97,6 +98,7 @@ cdef class _DynamicStructBuilder: cpdef _which_str(self) cpdef adopt(self, field, _DynamicOrphan orphan) cpdef disown(self, field) + cpdef get_data_as_view(self, field) cpdef as_reader(self) cpdef copy(self, num_first_segment_words=?, allocate_seg_callable=?) diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 0d6f542..6449a99 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -15,8 +15,8 @@ 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, PyBUF_CONTIG_RO -from cpython.memoryview cimport PyMemoryView_FromMemory +from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO, PyBuffer_FillInfo +from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer from cpython.bytes cimport PyBytes_FromStringAndSize from cpython.exc cimport PyErr_Clear from cython.operator cimport dereference as deref @@ -669,7 +669,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 PyMemoryView_FromMemory( temp_data.begin(), temp_data.size(), PyBUF_READ) + return ((temp_data.begin())[:temp_data.size()]) elif type == capnp.TYPE_LIST: return _DynamicListReader()._init(self.asList(), parent) elif type == capnp.TYPE_STRUCT: @@ -703,7 +703,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 PyMemoryView_FromMemory( temp_data.begin(), temp_data.size(), PyBUF_WRITE) + return ((temp_data.begin())[:temp_data.size()]) elif type == capnp.TYPE_LIST: return _DynamicListBuilder()._init(self.asList(), parent) elif type == capnp.TYPE_STRUCT: @@ -1226,6 +1226,29 @@ cdef class _DynamicStructReader: cpdef _has_by_field(self, _StructSchemaField field): return self.thisptr.hasByField(field.thisptr) + cpdef get_data_as_view(self, field): + """ + Efficiently get a read-only memoryview for a DATA field without copying. + """ + cdef C_DynamicValue.Reader val + cdef capnp.Data.Reader temp_data + + try: + val = self.thisptr.get(field) + except KjException as e: + raise e._to_python() from None + + if val.getType() != capnp.TYPE_DATA: + raise TypeError("Field '{}' is not a DATA field".format(field)) + + temp_data = val.asData() + + # Return read-only memoryview + cdef Py_buffer buf + if PyBuffer_FillInfo(&buf, self, temp_data.begin(), temp_data.size(), 1, PyBUF_CONTIG_RO) < 0: + raise KjException("Failed to create buffer info") + return PyMemoryView_FromBuffer(&buf) + cpdef _which_str(self): try: return helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr() @@ -1628,6 +1651,32 @@ cdef class _DynamicStructBuilder: """ return _DynamicOrphan()._init(self.thisptr.disown(field), self._parent) + cpdef get_data_as_view(self, field): + """ + Efficiently get a writable memoryview for a DATA field without copying. + + This allows in-place modification of the underlying buffer: + msg.get_data_as_view('myField')[0] = 0xFF + """ + cdef C_DynamicValue.Builder val + cdef capnp.Data.Builder temp_data + + try: + val = self.thisptr.get(field) + except KjException as e: + raise e._to_python() from None + + if val.getType() != capnp.TYPE_DATA: + raise TypeError("Field '{}' is not a DATA field".format(field)) + + temp_data = val.asData() + + # Return writable memoryview + cdef Py_buffer buf + if PyBuffer_FillInfo(&buf, self, temp_data.begin(), temp_data.size(), 0, PyBUF_WRITABLE) < 0: + raise KjException("Failed to create buffer info") + return PyMemoryView_FromBuffer(&buf) + cpdef as_reader(self): """A method for casting this Builder to a Reader diff --git a/test/test_get_data_view.py b/test/test_get_data_view.py new file mode 100644 index 0000000..cc20a60 --- /dev/null +++ b/test/test_get_data_view.py @@ -0,0 +1,210 @@ +import os +import pytest +import capnp +import sys +import gc + + +@pytest.fixture(scope="module") +def all_types(): + """Load the standard all_types.capnp schema.""" + directory = os.path.dirname(__file__) + return capnp.load(os.path.join(directory, "all_types.capnp")) + + +def test_set_bytes_get_bytes(all_types): + """ + Scenario 1: Set Byte -> Get Byte + Verify standard behavior: writing bytes results in reading bytes. + """ + msg = all_types.TestAllTypes.new_message() + input_data = b"hello_world" + + # Set + msg.dataField = input_data + + # Get + output_data = msg.dataField + + # Verify + assert isinstance(output_data, bytes) + assert output_data == input_data + + +def test_set_view_get_bytes(all_types): + """ + Scenario 2: Set View -> Get Byte + Verify compatibility: Passing a memoryview sets the data, + but standard attribute access returns a bytes copy. + """ + msg = all_types.TestAllTypes.new_message() + + # Create a memoryview source + raw_source = bytearray(b"view_source") + view = memoryview(raw_source) + + # Set via memoryview + msg.dataField = view + + # Get via standard attribute + output_data = msg.dataField + + # Verify + assert isinstance(output_data, bytes) + assert output_data == b"view_source" + + +def test_set_bytes_get_view_and_modify(all_types): + """ + Scenario 3: Set Byte -> Get View + Verify the high-performance API get_data_as_view. + The view must be writable and modifications must reflect in the message. + """ + msg = all_types.TestAllTypes.new_message() + + # Initial write + msg.dataField = b"ABCDE" + + # Get view via new API + view = msg.get_data_as_view("dataField") + + # Verify view properties + assert isinstance(view, memoryview) + assert view.readonly is False + assert view.tobytes() == b"ABCDE" + + # Verify in-place modification + view[0] = ord("Z") # Change 'A' to 'Z' + + # Verify modification is reflected in standard access + assert msg.dataField == b"ZBCDE" + + +def test_reader_vs_builder_view(all_types): + """ + Verify that Builder views are writable, but Reader views are read-only. + """ + # 1. Builder phase + builder = all_types.TestAllTypes.new_message() + builder.dataField = b"test_rw" + + builder_view = builder.get_data_as_view("dataField") + assert builder_view.readonly is False + builder_view[0] = ord("T") # Modification allowed + + # 2. Reader phase + reader = builder.as_reader() + + # Standard Get + assert reader.dataField == b"Test_rw" + + # Reader get_data_as_view + reader_view = reader.get_data_as_view("dataField") + assert isinstance(reader_view, memoryview) + assert reader_view.readonly is True + + # Attempting to modify Reader view should raise TypeError + with pytest.raises(TypeError): + reader_view[0] = ord("X") + + +def test_nested_struct_data(all_types): + """ + Verify that get_data_as_view works correctly on nested structs. + """ + msg = all_types.TestAllTypes.new_message() + + # Initialize nested struct + inner = msg.init("structField") + inner.int32Field = 100 + inner.dataField = b"nested_data" + + # 1. Verify standard access + assert msg.structField.dataField == b"nested_data" + + # 2. Verify nested get_data_as_view + view = msg.structField.get_data_as_view("dataField") + + assert isinstance(view, memoryview) + assert view.tobytes() == b"nested_data" + + # Modify nested data + view[0] = ord("N") + assert msg.structField.dataField == b"Nested_data" + + +def test_corner_cases_values(all_types): + """ + Test edge cases: Empty bytes and binary data with nulls. + """ + msg = all_types.TestAllTypes.new_message() + + # Case A: Empty Bytes + msg.dataField = b"" + assert msg.dataField == b"" + view = msg.get_data_as_view("dataField") + assert len(view) == 0 + + # Case B: Binary data containing null bytes + binary_data = b"\x00\xff\x00\x01" + msg.dataField = binary_data + assert msg.dataField == binary_data + assert msg.get_data_as_view("dataField").tobytes() == binary_data + + +def test_error_wrong_type(all_types): + """ + Test error handling: Calling get_data_as_view on non-Data fields. + """ + msg = all_types.TestAllTypes.new_message() + msg.int32Field = 123 + msg.textField = "I am text" + + # Attempt on Int field + with pytest.raises(TypeError) as excinfo: + msg.get_data_as_view("int32Field") + assert "not a DATA field" in str(excinfo.value) + + # Attempt on Text field + with pytest.raises(TypeError) as excinfo: + msg.get_data_as_view("textField") + assert "not a DATA field" in str(excinfo.value) + + +def test_error_missing_field(all_types): + """ + Test error handling: Accessing a non-existent field name. + """ + msg = all_types.TestAllTypes.new_message() + + # Accessing a missing field should raise AttributeError (standard Python behavior) + with pytest.raises(AttributeError) as excinfo: + msg.get_data_as_view("non_existent_field") + + # Optional: Verify the error message contains the field name + assert "non_existent_field" in str(excinfo.value) + + +def test_view_keeps_message_alive(all_types): + """ + Verify that a View keeps messages alive. + """ + msg = all_types.TestAllTypes.new_message() + expected_data = b"persistence_check" + msg.dataField = expected_data + + initial_ref_count = sys.getrefcount(msg) + view = msg.get_data_as_view("dataField") + new_ref_count = sys.getrefcount(msg) + + assert ( + new_ref_count > initial_ref_count + ), f"View failed to hold reference to Message! (Old: {initial_ref_count}, New: {new_ref_count})" + print( + f"\n[Ref Check] Success: Ref count increased from {initial_ref_count} to {new_ref_count}" + ) + + del msg + gc.collect() + + assert view.tobytes() == expected_data diff --git a/test/test_py_custom_message_builder.py b/test/test_py_custom_message_builder.py index 1b62e1a..1b5b8f1 100644 --- a/test/test_py_custom_message_builder.py +++ b/test/test_py_custom_message_builder.py @@ -41,9 +41,6 @@ def test_addressbook(all_types): 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"