diff --git a/opendis/record.py b/opendis/record.py index 7c8b082..211cd8a 100644 --- a/opendis/record.py +++ b/opendis/record.py @@ -4,24 +4,52 @@ """ from collections.abc import Sequence -from ctypes import _SimpleCData, BigEndianStructure, c_uint +from ctypes import ( + _SimpleCData, + BigEndianStructure, + c_uint8, + c_uint16, + c_uint32, + sizeof, +) +from typing import Literal + +from .stream import DataInputStream, DataOutputStream from .types import ( bf_enum, bf_int, bf_uint, ) -from .DataInputStream import DataInputStream -from .DataOutputStream import DataOutputStream - - -def _bitfield( - name: str, - bytesize: int, - fields: Sequence[ - tuple[str, type[_SimpleCData]] | tuple[str, type[_SimpleCData], int] - ], - ): +# Type definitions for bitfield field descriptors +CTypeFieldDescription = tuple[str, type[_SimpleCData], int] +DisFieldDescription = tuple[str, "DisFieldType", int] + +# Field type constants simplify the construction of bitfields +# which would otherwise require manually specifying ctypes types. +# The currently implemented bitfields only use integers, but DIS7 +# mentions CHAR types which may be needed in future. +DisFieldType = Literal["INTEGER"] +INTEGER = "INTEGER" + + +def field(name: str, + ftype: DisFieldType, + bits: int) -> CTypeFieldDescription: + """Helper function to create the field description tuple used by ctypes.""" + match (ftype, bits): + case (INTEGER, b) if 0 < b <= 8: + return (name, c_uint8, bits) + case (INTEGER, b) if 8 < b <= 16: + return (name, c_uint16, bits) + case (INTEGER, b) if 16 < b <= 32: + return (name, c_uint32, bits) + case _: + raise ValueError(f"Unrecognized (ftype, bits): {ftype}, {bits}") + + +def _bitfield(name: str, + fields: Sequence[DisFieldDescription]): """Factory function for bitfield structs, which are subclasses of ctypes.Structure. These are used in records that require them to unpack non-octet-sized fields. @@ -29,15 +57,32 @@ def _bitfield( Args: name: Name of the bitfield struct. bytesize: Size of the bitfield in bytes. - fields: Sequence of tuples defining the fields of the bitfield. - See https://docs.python.org/3/library/ctypes.html#ctypes.Structure._fields_ + fields: Sequence of tuples defining fields of the bitfield, in the form + (field_name, "INTEGER", field_size_in_bits). """ - if bytesize <= 0: - raise ValueError("Cannot create bitfield with less than one byte") - + # Argument validation + struct_fields = [] + bitsize = 0 + for name, ftype, bits in fields: + if ftype not in (INTEGER,): + raise ValueError(f"Unsupported field type: {ftype}") + if not isinstance(bits, int): + raise ValueError(f"Field size must be int: {bits!r}") + if bits <= 0 or bits > 32: + raise ValueError(f"Field size must be between 1 and 32: got {bits}") + bitsize += bits + struct_fields.append(field(name, ftype, bits)) + + if bitsize == 0: + raise ValueError(f"Bitfield size cannot be zero") + elif bitsize % 8 != 0: + raise ValueError(f"Bitfield size must be multiple of 8, got {bitsize}") + bytesize = bitsize // 8 + + # Create the struct class class Bitfield(BigEndianStructure): - _fields_ = fields - + _fields_ = struct_fields + @staticmethod def marshalledSize() -> int: return bytesize @@ -48,6 +93,12 @@ def serialize(self, outputStream: DataOutputStream) -> None: @classmethod def parse(cls, inputStream: DataInputStream) -> "Bitfield": return cls.from_buffer_copy(inputStream.read_bytes(bytesize)) + + # Sanity check: ensure the struct size matches expected size + assert sizeof(Bitfield) == bytesize, \ + f"Bitfield size mismatch: expected {bytesize}, got {sizeof(Bitfield)}" + + # Assign the class name Bitfield.__name__ = name return Bitfield @@ -61,11 +112,11 @@ class NetId: YY = Frequency Table """ - _struct = _bitfield(name="NetId", bytesize=2, fields=[ - ("netNumber", c_uint, 10), - ("frequencyTable", c_uint, 2), - ("mode", c_uint, 2), - ("padding", c_uint, 2) + _struct = _bitfield(name="NetId", fields=[ + ("netNumber", INTEGER, 10), + ("frequencyTable", INTEGER, 2), + ("mode", INTEGER, 2), + ("padding", INTEGER, 2) ]) def __init__(self, @@ -113,11 +164,11 @@ class SpreadSpectrum: In Python, the presence or absence of each technique is indicated by a bool. """ - _struct = _bitfield("SpreadSpectrum", 2, [ - ("frequencyHopping", c_uint, 1), - ("pseudoNoise", c_uint, 1), - ("timeHopping", c_uint, 1), - ("padding", c_uint, 13) + _struct = _bitfield(name="SpreadSpectrum", fields=[ + ("frequencyHopping", INTEGER, 1), + ("pseudoNoise", INTEGER, 1), + ("timeHopping", INTEGER, 1), + ("padding", INTEGER, 13) ]) def __init__(self,