Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions capnp/lib/capnp.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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=?)

Expand Down Expand Up @@ -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=?)
Expand Down
57 changes: 53 additions & 4 deletions capnp/lib/capnp.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -669,7 +669,7 @@ cdef to_python_reader(C_DynamicValue.Reader self, object parent):
return (<char*>temp_text.begin())[:temp_text.size()]
elif type == capnp.TYPE_DATA:
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:
return _DynamicListReader()._init(self.asList(), parent)
elif type == capnp.TYPE_STRUCT:
Expand Down Expand Up @@ -703,7 +703,7 @@ cdef to_python_builder(C_DynamicValue.Builder self, object parent):
return (<char*>temp_text.begin())[:temp_text.size()]
elif type == capnp.TYPE_DATA:
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:
return _DynamicListBuilder()._init(self.asList(), parent)
elif type == capnp.TYPE_STRUCT:
Expand Down Expand Up @@ -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, <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):
try:
return <char *>helpers.fixMaybe(self.thisptr.which()).getProto().getName().cStr()
Expand Down Expand Up @@ -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, <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):
"""A method for casting this Builder to a Reader

Expand Down
210 changes: 210 additions & 0 deletions test/test_get_data_view.py
Original file line number Diff line number Diff line change
@@ -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
3 changes: 0 additions & 3 deletions test/test_py_custom_message_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,9 +41,6 @@ def __call__(self, minimum_size: int) -> bytearray:
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"

Expand Down
Loading