diff --git a/src/packages/harp-protocol/src/harp/protocol/_constants.py b/src/packages/harp-protocol/src/harp/protocol/_constants.py index e2bdc0c..6b23e88 100644 --- a/src/packages/harp-protocol/src/harp/protocol/_constants.py +++ b/src/packages/harp-protocol/src/harp/protocol/_constants.py @@ -12,6 +12,9 @@ _HEADER_LEN: int = 5 """Fixed header size in bytes: msg_type + length + address + port + payload_type.""" +_MIN_FRAME_LEN: int = 6 +"""Smallest frame on the wire in bytes, the fixed header plus the checksum.""" + _TIMESTAMP_LEN: int = 6 """Timestamp field size in bytes: 4-byte seconds as u32 plus 2-byte microseconds as u16.""" diff --git a/src/packages/harp-protocol/src/harp/protocol/_message.py b/src/packages/harp-protocol/src/harp/protocol/_message.py index 7578ad5..020a517 100644 --- a/src/packages/harp-protocol/src/harp/protocol/_message.py +++ b/src/packages/harp-protocol/src/harp/protocol/_message.py @@ -11,6 +11,7 @@ from ._constants import ( _DEFAULT_PORT, _HEADER_LEN, + _MIN_FRAME_LEN, _TICK_PERIOD_S, _TIMESTAMP_FLAG, _TIMESTAMP_LEN, @@ -81,8 +82,8 @@ def parse(cls, data: bytes | bytearray | memoryview) -> "HarpMessage[Any]": """Parse and validate a complete Harp Message from a byte sequence. Raises ``HarpParseError`` on failure.""" raw = data if isinstance(data, bytes) else bytes(data) - if len(raw) < 6: - raise HarpParseError(f"Frame too short: {len(raw)} bytes (minimum 6)") + if len(raw) < _MIN_FRAME_LEN: + raise HarpParseError(f"Frame too short: {len(raw)} bytes (minimum {_MIN_FRAME_LEN})") if not _validate_checksum(raw): raise HarpParseError("Checksum mismatch") diff --git a/src/packages/harp-protocol/src/harp/protocol/_payload.py b/src/packages/harp-protocol/src/harp/protocol/_payload.py index 1f14871..2bb7095 100644 --- a/src/packages/harp-protocol/src/harp/protocol/_payload.py +++ b/src/packages/harp-protocol/src/harp/protocol/_payload.py @@ -1030,7 +1030,18 @@ def __init__(self, value: object = _MISSING_INIT, /, **kwargs: object) -> None: raise TypeError(f"{type(self).__name__}() requires a value") if kwargs: raise TypeError(f"{type(self).__name__}() got unexpected kwargs: {sorted(kwargs)}") - self._arr = np.asarray(value, dtype=self.payload_dtype) + subdtype = self.payload_dtype.subdtype + if subdtype is not None: + element_dtype, shape = subdtype + arr = np.asarray(value, dtype=element_dtype) + if arr.shape != shape: + expected = shape[0] if len(shape) == 1 else shape + raise ValueError( + f"{type(self).__name__}() expects {expected} elements but got {arr.size}" + ) + else: + arr = np.asarray(value, dtype=self.payload_dtype) + self._arr = arr @classmethod def _unwrap(cls, arr: "np.ndarray") -> Any: diff --git a/src/packages/harp-protocol/src/harp/protocol/_register.py b/src/packages/harp-protocol/src/harp/protocol/_register.py index effcfe0..31109e1 100644 --- a/src/packages/harp-protocol/src/harp/protocol/_register.py +++ b/src/packages/harp-protocol/src/harp/protocol/_register.py @@ -9,12 +9,13 @@ from ._constants import ( _DEFAULT_PORT, _HEADER_LEN, + _MIN_FRAME_LEN, _TICK_PERIOD_S, _TIMESTAMP_FLAG, _TIMESTAMPED_PAYLOAD_OFFSET, _TS_MICROS_OFFSET, ) -from ._message import HarpMessage +from ._message import HarpMessage, HarpParseError from ._message_type import MessageType, message_type_to_byte from ._payload import ( Batch, @@ -165,6 +166,12 @@ def parse(cls, value: HarpMessage | bytes | bytearray | memoryview) -> U: registers) return the raw numpy scalar or ndarray directly. """ buf = value.payload_bytes if isinstance(value, HarpMessage) else value + expected = cls.payload_class.payload_dtype.itemsize + if len(buf) < expected: + raise HarpParseError( + f"{cls.__name__} reads {expected} payload bytes as {cls.payload_type!r} " + f"but only {len(buf)} are available." + ) record = np.frombuffer(buf, dtype=cls.payload_class.payload_dtype, count=1)[0] return cast(U, cls.payload_class._unwrap(record)) @@ -185,9 +192,20 @@ def parse_bulk( payload = payload_cls._from_array(np.empty(0, dtype=payload_cls.payload_dtype)) return data, None, None, cast("Batch[Any]", payload) + if len(data) < _MIN_FRAME_LEN: + raise HarpParseError( + f"{cls.__name__} reads frames of at least {_MIN_FRAME_LEN} bytes " + f"but only {len(data)} are available." + ) + stride = ( int(data[1]) + 2 ) # TODO this assumes all frames have the same length but we may want to revisit in the future. + if len(data) < stride: + raise HarpParseError( + f"{cls.__name__} reads frames of {stride} bytes as declared by the first " + f"length byte, but only {len(data)} are available." + ) nrows = len(data) // stride is_timestamped = bool(int(data[4]) & _TIMESTAMP_FLAG) payload_offset = _TIMESTAMPED_PAYLOAD_OFFSET if is_timestamped else _HEADER_LEN @@ -235,22 +253,35 @@ def format_bulk( ``parse_bulk``. """ payload_cls = cls.payload_class - itemsize = payload_cls.payload_dtype.itemsize + record_dtype = payload_cls.payload_dtype + itemsize = record_dtype.itemsize + subdtype = record_dtype.subdtype if isinstance(values, PayloadBase): records = np.atleast_1d(np.asarray(values.payload_array)) else: records = np.atleast_1d(np.asarray(values)) + if subdtype is not None and records.dtype != record_dtype: + # An array register: convert against the element type and let the declared + # shape decide the frame count, as `format` does for a single frame. + element_dtype, shape = subdtype + records = np.asarray(records, dtype=element_dtype) + if records.shape[-len(shape) :] != shape: + raise ValueError( + f"{cls.__name__}.format_bulk: values of shape {records.shape} do not end " + f"in the declared payload shape {shape}" + ) + records = records.reshape(-1, *shape) + elif ( # Coerce the element type only for plain scalar payloads (e.g. an int - # list for a scalar register). Struct/sub-array records already carry - # the right byte layout and must not be re-cast. - plain = ( - records.dtype.names is None - and records.dtype.subdtype is None - and payload_cls.payload_dtype.names is None - and payload_cls.payload_dtype.subdtype is None - ) - if plain and records.dtype != payload_cls.payload_dtype: - records = records.astype(payload_cls.payload_dtype) + # list for a scalar register). Struct records already carry the right + # byte layout and must not be re-cast. + records.dtype.names is None + and records.dtype.subdtype is None + and record_dtype.names is None + and subdtype is None + and records.dtype != record_dtype + ): + records = records.astype(record_dtype) nrows = len(records) flat = np.ascontiguousarray(records).tobytes() if len(flat) != nrows * itemsize: diff --git a/tests/protocol/test_register.py b/tests/protocol/test_register.py index a417926..e7357d1 100644 --- a/tests/protocol/test_register.py +++ b/tests/protocol/test_register.py @@ -5,7 +5,7 @@ import numpy as np import pytest from harp.data import parse_to_dataframe, payload_to_dataframe, to_buffer, to_file -from harp.protocol._message import HarpMessage +from harp.protocol._message import HarpMessage, HarpParseError from harp.protocol._message_type import MessageType from harp.protocol._payload import ( PayloadBase, @@ -32,6 +32,7 @@ RegisterS64, RegisterU8, RegisterU16, + RegisterU16Array, RegisterU32, RegisterU32Array, RegisterU64, @@ -196,6 +197,32 @@ def test_format_with_payload_instance(reg_cls, payload_cls, value): assert msg.payload_bytes == payload.payload_array.tobytes() +def test_format_accepts_sequence_for_array_register(): + # The payload dtype is a sub-array, so converting a sequence against it directly + # broadcasts each element into the full shape and doubles the payload. + reg = RegisterU16Array(0x20, length=2) + expected = np.array([1, 2], dtype=np.uint16).tobytes() + for value in ([1, 2], (1, 2), np.array([1, 2], dtype=np.uint16)): + msg = _parse_frame(reg.format(value)) + assert msg.payload_bytes == expected + assert list(reg.parse(msg)) == [1, 2] + + +def test_format_rejects_wrong_length_sequence(): + reg = RegisterU16Array(0x20, length=2) + with pytest.raises(ValueError, match="expects 2 elements but got 3"): + reg.format([1, 2, 3]) + + +def test_parse_names_register_on_short_payload(): + # A read request carries no payload, and numpy would otherwise report only + # "buffer is smaller than requested size", naming neither side. + reg = RegisterU16(0x20) + request = _parse_frame(reg.format(message_type=MessageType.Read)) + with pytest.raises(HarpParseError, match="reads 2 payload bytes"): + reg.parse(request) + + def test_format_with_payload_instance_via_register(): """format() accepts a typed PayloadU32 and encodes it correctly.""" payload = PayloadU32(42) @@ -579,6 +606,56 @@ def test_format_bulk_is_exact_inverse_of_parse_bulk(): assert bytes(rebuilt) == bytes(original) +def test_format_bulk_accepts_sequences_for_array_register(): + # An array register formats a sequence element-wise, as format does, so every + # spelling of the same values produces the same frames. + reg = RegisterU16Array(0x20, length=2) + records = np.array([(1, 2), (3, 4)], dtype=reg.payload_class.payload_dtype) + expected = bytes(reg.format_bulk(records)) + for values in ( + np.array([[1, 2], [3, 4]], dtype=np.uint16), + np.array([[1, 2], [3, 4]]), + [[1, 2], [3, 4]], + ((1, 2), (3, 4)), + ): + assert bytes(reg.format_bulk(values)) == expected + + +def test_format_bulk_single_frame_matches_format_for_array_register(): + reg = RegisterU16Array(0x20, length=2) + bulk = reg.format_bulk([1, 2], message_type=MessageType.Write) + assert bytes(bulk) == reg.format([1, 2], message_type=MessageType.Write) + + +def test_format_bulk_rejects_shape_without_declared_payload(): + # Four elements for a two-element register could be one frame or two, so the + # ambiguous flat form raises rather than guessing. + reg = RegisterU16Array(0x20, length=2) + with pytest.raises(ValueError, match="do not end in the declared payload shape"): + reg.format_bulk(np.array([1, 2, 3, 4], dtype=np.uint16)) + with pytest.raises(ValueError, match="do not end in the declared payload shape"): + reg.format_bulk([[1, 2, 3], [4, 5, 6]]) + + +def test_parse_bulk_names_register_on_partial_frame(): + # numpy would otherwise report an out-of-bounds index against the strided view, + # naming neither the register nor how many bytes a frame needs. + reg = RegisterU16Array(0x20, length=2) + buf = bytes(reg.format_bulk([[1, 2], [3, 4]])) + stride = buf[1] + 2 + with pytest.raises(HarpParseError, match="at least 6 bytes"): + reg.parse_bulk(buf[:5]) + with pytest.raises(HarpParseError, match=f"frames of {stride} bytes"): + reg.parse_bulk(buf[: stride - 1]) + + +def test_parse_bulk_empty_buffer_does_not_raise(): + reg = RegisterU16Array(0x20, length=2) + _data, timestamps, msgtype, payload = reg.parse_bulk(b"") + assert len(np.asarray(payload.payload_array)) == 0 + assert timestamps is None and msgtype is None + + def test_to_buffer_and_to_file_roundtrip(tmp_path): reg = RegisterU16(0x20) values = np.array([10, 20], dtype="