import struct
from typing import (
Any,
Self,
ClassVar,
)
from abc import ABC
from dataclasses import (
field as datafield,
dataclass,
)
__all__ = (
'LayoutField',
'LayoutMixin',
)
[docs]
@dataclass(frozen=True, slots=True)
class LayoutField:
"""
A field in a byte-level layout.
Attributes:
name: the name of the field
format: the format of the field used by the `struct` module, excluding the endian prefix
size: the size of the field, automatically calculated from its format
"""
name: str
format: str
size: int = datafield(init=False)
def __post_init__(self) -> None:
object.__setattr__(self, 'size', struct.calcsize(self.format))
[docs]
class LayoutMixin(ABC):
"""
An abstract base class for defining byte-level layouts.
It holds a buffer internally and manages its layout according to the configured `fields`.
Attributes:
buffer: the underlying buffer organized according to the layout
Raises:
TypeError: if `fields` is not properly defined
"""
ENDIAN: ClassVar[str] = '<'
"""The endian format used by the `struct` module."""
size: ClassVar[int]
"""The total size of the layout, automatically calculated from `fields`."""
lookup: ClassVar[dict[str, tuple[int, LayoutField]]]
"""A mapping from each field name to its offset and field, generated automatically."""
fields: ClassVar[tuple[LayoutField, ...]]
"""The fields that define the layout."""
buffer: bytearray
def __init_subclass__(cls: type[Self]) -> None: # pragma: no cover
cname = cls.__name__
if not hasattr(cls, 'fields'):
raise TypeError(f"{cname} must define the class attribute fields.")
if not isinstance(cls.fields, tuple):
raise TypeError(f"The class attribute fields of {cname} must be of type tuple.")
for field in cls.fields:
if not isinstance(field, LayoutField):
raise TypeError(f"The class attribute fields of {cname} must contain only LayoutField objects.")
cls.lookup = {}
offset = 0
for field in cls.fields:
cls.lookup[field.name] = (offset, field)
offset += field.size
cls.size = offset
def __init__(self, buffer: bytearray) -> None:
assert len(buffer) == self.size
self.buffer = buffer
[docs]
def get(self, name: str) -> Any:
"""
Return the value of the field with the given name from the buffer.
"""
offset, field = self.lookup[name]
format = self.get_format(field)
return struct.unpack_from(format, self.buffer, offset)[0]
[docs]
def set(self, name: str, value: Any) -> None:
"""
Set the value of the field with the given name in the buffer.
"""
offset, field = self.lookup[name]
format = self.get_format(field)
struct.pack_into(format, self.buffer, offset, value)