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
This commit is contained in:
@@ -65,6 +65,7 @@ cdef class _DynamicStructReader:
|
|||||||
cpdef _which_str(self)
|
cpdef _which_str(self)
|
||||||
cpdef _get_by_field(self, _StructSchemaField field)
|
cpdef _get_by_field(self, _StructSchemaField field)
|
||||||
cpdef _has_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=?)
|
cpdef as_builder(self, num_first_segment_words=?, allocate_seg_callable=?)
|
||||||
|
|
||||||
@@ -97,6 +98,7 @@ cdef class _DynamicStructBuilder:
|
|||||||
cpdef _which_str(self)
|
cpdef _which_str(self)
|
||||||
cpdef adopt(self, field, _DynamicOrphan orphan)
|
cpdef adopt(self, field, _DynamicOrphan orphan)
|
||||||
cpdef disown(self, field)
|
cpdef disown(self, field)
|
||||||
|
cpdef get_data_as_view(self, field)
|
||||||
|
|
||||||
cpdef as_reader(self)
|
cpdef as_reader(self)
|
||||||
cpdef copy(self, num_first_segment_words=?, allocate_seg_callable=?)
|
cpdef copy(self, num_first_segment_words=?, allocate_seg_callable=?)
|
||||||
|
|||||||
@@ -15,8 +15,8 @@ from capnp.includes.schema_cpp cimport (MessageReader,)
|
|||||||
|
|
||||||
from builtins import memoryview as BuiltinsMemoryview
|
from builtins import memoryview as BuiltinsMemoryview
|
||||||
from cpython cimport array, Py_buffer, PyObject_CheckBuffer
|
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.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO, PyBuffer_FillInfo
|
||||||
from cpython.memoryview cimport PyMemoryView_FromMemory
|
from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer
|
||||||
from cpython.bytes cimport PyBytes_FromStringAndSize
|
from cpython.bytes cimport PyBytes_FromStringAndSize
|
||||||
from cpython.exc cimport PyErr_Clear
|
from cpython.exc cimport PyErr_Clear
|
||||||
from cython.operator cimport dereference as deref
|
from cython.operator cimport dereference as deref
|
||||||
@@ -669,7 +669,7 @@ cdef to_python_reader(C_DynamicValue.Reader self, object parent):
|
|||||||
return (<char*>temp_text.begin())[:temp_text.size()]
|
return (<char*>temp_text.begin())[:temp_text.size()]
|
||||||
elif type == capnp.TYPE_DATA:
|
elif type == capnp.TYPE_DATA:
|
||||||
temp_data = self.asData()
|
temp_data = self.asData()
|
||||||
return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_READ)
|
return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
|
||||||
elif type == capnp.TYPE_LIST:
|
elif type == capnp.TYPE_LIST:
|
||||||
return _DynamicListReader()._init(self.asList(), parent)
|
return _DynamicListReader()._init(self.asList(), parent)
|
||||||
elif type == capnp.TYPE_STRUCT:
|
elif type == capnp.TYPE_STRUCT:
|
||||||
@@ -703,7 +703,7 @@ cdef to_python_builder(C_DynamicValue.Builder self, object parent):
|
|||||||
return (<char*>temp_text.begin())[:temp_text.size()]
|
return (<char*>temp_text.begin())[:temp_text.size()]
|
||||||
elif type == capnp.TYPE_DATA:
|
elif type == capnp.TYPE_DATA:
|
||||||
temp_data = self.asData()
|
temp_data = self.asData()
|
||||||
return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_WRITE)
|
return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
|
||||||
elif type == capnp.TYPE_LIST:
|
elif type == capnp.TYPE_LIST:
|
||||||
return _DynamicListBuilder()._init(self.asList(), parent)
|
return _DynamicListBuilder()._init(self.asList(), parent)
|
||||||
elif type == capnp.TYPE_STRUCT:
|
elif type == capnp.TYPE_STRUCT:
|
||||||
@@ -1226,6 +1226,29 @@ cdef class _DynamicStructReader:
|
|||||||
cpdef _has_by_field(self, _StructSchemaField field):
|
cpdef _has_by_field(self, _StructSchemaField field):
|
||||||
return self.thisptr.hasByField(field.thisptr)
|
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, <void*>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):
|
cpdef _which_str(self):
|
||||||
try:
|
try:
|
||||||
return <char *>helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr()
|
return <char *>helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr()
|
||||||
@@ -1628,6 +1651,32 @@ cdef class _DynamicStructBuilder:
|
|||||||
"""
|
"""
|
||||||
return _DynamicOrphan()._init(self.thisptr.disown(field), self._parent)
|
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, <void*>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):
|
cpdef as_reader(self):
|
||||||
"""A method for casting this Builder to a Reader
|
"""A method for casting this Builder to a Reader
|
||||||
|
|
||||||
|
|||||||
210
test/test_get_data_view.py
Normal file
210
test/test_get_data_view.py
Normal file
@@ -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
|
||||||
@@ -41,9 +41,6 @@ def test_addressbook(all_types):
|
|||||||
struct_builder.init("dataField", 5)
|
struct_builder.init("dataField", 5)
|
||||||
assert struct_builder._get("dataField") == b"\x00\x00\x00\x00\x00"
|
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"
|
struct_builder.dataField = b"hello"
|
||||||
assert struct_builder._get("dataField") == b"hello"
|
assert struct_builder._get("dataField") == b"hello"
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user