diff --git a/README.md b/README.md index 79a90c6..6592a4d 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,11 @@ s.decode([4, 2, 65]) # MyStruct(myNum=1026, myLetter=b"A") s.decode([4, 2, 65], little_endian=True) # MyStruct(myNum=516, myLetter=b"A") + +# We can modify the class values and encode the data to send back +s.myNum = 2562 +s.encode() +# [10, 2, 65] ``` For arrays of basic elements, you need to Annotate them with @@ -75,9 +80,10 @@ List defaults will set all items in the list to the same value. Currently setting a complete default list for all values is not implemented. ```c -struct MyStruct P +struct MyStruct { uint8_t myInt = 5; uint8_t myInts[2]; +}; ``` ```python diff --git a/src/pystructtype/__init__.py b/src/pystructtype/__init__.py index 9e2120b..df1f6ed 100644 --- a/src/pystructtype/__init__.py +++ b/src/pystructtype/__init__.py @@ -1,482 +1,43 @@ -import inspect -import itertools -import re -import struct -from collections.abc import Callable, Generator -from copy import deepcopy -from dataclasses import dataclass, field, is_dataclass -from typing import ( - Annotated, - Any, - TypeVar, - cast, - get_args, - get_origin, - get_type_hints, - overload, +from pystructtype.bitstype import BitsType, bits +from pystructtype.structdataclass import StructDataclass, struct_dataclass +from pystructtype.structtypes import ( + TypeInfo, + TypeMeta, + char_t, + double_t, + float_t, + int8_t, + int16_t, + int32_t, + int64_t, + uint8_t, + uint16_t, + uint32_t, + uint64_t, ) - -def list_chunks(_list: list, n: int) -> Generator[list]: - """ - Yield successive n-sized chunks from a list. - :param _list: List to chunk out - :param n: Size of chunks - :return: Generator of n-sized chunks of _list - """ - yield from (_list[i : i + n] for i in range(0, len(_list), n)) - - -def type_from_annotation(_type: type) -> type: - """ - Find the base type from an Annotated type, or return it unchanged if - not Annotated - :param _type: Type to check - :return: Annotated base type or the given type if not Annotated - """ - # If we have an origin for the given type, and it's Annotated - if (origin := get_origin(_type)) and origin is Annotated: - # Keep running `get_args` on the first element of whatever - # `get_args` returns, until we get nothing back - arg = _type - t: Any = _type - while t := get_args(t): - arg = t[0] - - # This will be the base type - return arg - # No origin, or the origin is not Annotated, just return the given type - return _type - - -T = TypeVar("T", int, float, default=int) - - -@dataclass(frozen=True) -class TypeMeta[T]: - size: int = 1 - default: T | None = None - - -@dataclass(frozen=True) -class TypeInfo: - format: str - byte_size: int - - -# TODO: Support proper "c-string" types - -# Fixed Size Types -char_t = Annotated[int, TypeInfo("c", 1)] -int8_t = Annotated[int, TypeInfo("b", 1)] -uint8_t = Annotated[int, TypeInfo("B", 1)] -int16_t = Annotated[int, TypeInfo("h", 2)] -uint16_t = Annotated[int, TypeInfo("H", 2)] -int32_t = Annotated[int, TypeInfo("i", 4)] -uint32_t = Annotated[int, TypeInfo("I", 4)] -int64_t = Annotated[int, TypeInfo("q", 8)] -uint64_t = Annotated[int, TypeInfo("Q", 8)] - -# TODO: Make a special Bool class to auto-convert from int to bool - -# Named Types -float_t = Annotated[float, TypeInfo("f", 4)] -double_t = Annotated[float, TypeInfo("d", 8)] - - -@dataclass -class TypeIterator: - key: str - base_type: type - type_info: TypeInfo | None - type_meta: TypeMeta | None - is_list: bool - is_pystructtype: bool - - @property - def size(self): - return getattr(self.type_meta, "size", 1) - - -def iterate_types(cls) -> Generator[TypeIterator]: - for key, hint in get_type_hints(cls, include_extras=True).items(): - # Grab the base type from a possibly annotated type hint - base_type = type_from_annotation(hint) - - # Determine if the type is a list - # ex. list[bool] (yes) vs bool (no) - is_list = issubclass(origin, list) if (origin := get_origin(base_type)) else False - - # Grab the type hints top args and look for any TypeMeta objects - type_args = get_args(hint) - type_meta = next((x for x in type_args if isinstance(x, TypeMeta)), None) - - # type_args has the possibility of being nested within more tuples - # drill down the type_args until we hit empty, then we know we're at the bottom - # which is where type_info will exist - if type_args and len(type_args) > 1: - while args := get_args(type_args[0]): - type_args = args - - # Find the TypeInfo object on the lowest rung of the type_args - type_info = next((x for x in type_args if isinstance(x, TypeInfo)), None) - - # At this point we may have possibly drilled down into `type_args` to find the true base type - if type_args: - base_type = type_from_annotation(type_args[0]) - - # Determine if we are a subclass of a pystructtype - # If we have a type_info object in the Annotation, or we're actually a subtype of StructDataclass - is_pystructtype = type_info is not None or ( - inspect.isclass(base_type) and issubclass(base_type, StructDataclass) - ) - - yield TypeIterator(key, base_type, type_info, type_meta, is_list, is_pystructtype) - - -@dataclass -class StructState: - name: str - struct_fmt: str - size: int - - -class StructDataclass: - def __post_init__(self): - self._state: list[StructState] = [] - # Grab Struct Format - self.struct_fmt = "" - for type_iterator in iterate_types(self.__class__): - if type_iterator.type_info: - self._state.append( - StructState( - type_iterator.key, - type_iterator.type_info.format, - type_iterator.size, - ) - ) - self.struct_fmt += ( - f"{type_iterator.size if type_iterator.size > 1 else ''}{type_iterator.type_info.format}" - ) - elif inspect.isclass(type_iterator.base_type) and issubclass(type_iterator.base_type, StructDataclass): - attr = getattr(self, type_iterator.key) - if type_iterator.is_list: - fmt = attr[0].struct_fmt - else: - fmt = attr.struct_fmt - self._state.append(StructState(type_iterator.key, fmt, type_iterator.size)) - self.struct_fmt += fmt * type_iterator.size - else: - # We have no TypeInfo object, and we're not a StructDataclass - # This means we're a regularly defined class variable, and we - # Don't have to do anything about this. - pass - self._simplify_format() - self._byte_length = struct.calcsize("=" + self.struct_fmt) - # print(f"{self.__class__.__name__}: {self._byte_length} : {self.struct_fmt}") - - def _simplify_format(self) -> None: - # First expand the format - expanded_format = "" - items = re.findall(r"([a-zA-Z]|\d+)", self.struct_fmt) - items_len = len(items) - idx = 0 - while idx < items_len: - if "0" <= (item := items[idx]) <= "9": - idx += 1 - expanded_format += items[idx] * int(item) - else: - expanded_format += item - idx += 1 - - simplified_format = "" - for group in (x[0] for x in re.findall(r"(([a-zA-Z])\2*)", expanded_format)): - group_len = len(group) - simplified_format += f"{group_len if group_len > 1 else ''}{group[0]}" - - self.struct_fmt = simplified_format - - def size(self) -> int: - return sum(state.size for state in self._state) - - @staticmethod - def _endian(little_endian: bool) -> str: - return "<" if little_endian else ">" - - @staticmethod - def _to_bytes(data: list[int] | bytes) -> bytes: - if isinstance(data, bytes): - return data - return bytes(data) - - @staticmethod - def _to_list(data: list[int] | bytes) -> list[int]: - if isinstance(data, bytes): - return list(data) - return data - - def _decode(self, data: list[int]) -> None: - idx = 0 - - for state in self._state: - attr = getattr(self, state.name) - - if isinstance(attr, list) and isinstance(attr[0], StructDataclass): - list_idx = 0 - sub_struct_byte_length = attr[0].size() - while list_idx < state.size: - instance: StructDataclass = attr[list_idx] - instance._decode(data[idx : idx + sub_struct_byte_length]) - list_idx += 1 - idx += sub_struct_byte_length - elif isinstance(attr, StructDataclass): - if state.size != 1: - raise Exception("This should be a size of one, dingus") - - sub_struct_byte_length = attr.size() - attr._decode(data[idx : idx + sub_struct_byte_length]) - idx += sub_struct_byte_length - elif state.size == 1: - setattr(self, state.name, data[idx]) - idx += 1 - else: - list_idx = 0 - while list_idx < state.size: - getattr(self, state.name)[list_idx] = data[idx] - list_idx += 1 - idx += 1 - - def decode(self, data: list[int] | bytes, little_endian=False) -> None: - data = self._to_bytes(data) - - # Decode - self._decode(list(struct.unpack(self._endian(little_endian) + self.struct_fmt, data))) - - def _encode(self) -> list[int]: - result: list[int] = [] - - for state in self._state: - attr = getattr(self, state.name) - - if isinstance(attr, list) and isinstance(attr[0], StructDataclass): - item: StructDataclass - for item in attr: - result.extend(item._encode()) - elif isinstance(attr, StructDataclass): - if state.size != 1: - raise Exception("This should be a size of one, dingus") - result.extend(attr._encode()) - elif state.size == 1: - result.append(getattr(self, state.name)) - else: - result.extend(getattr(self, state.name)) - return result - - def encode(self, little_endian=False) -> bytes: - result = self._encode() - return struct.pack(self._endian(little_endian) + self.struct_fmt, *result) - - -@overload -def struct_dataclass(_cls: type[StructDataclass]) -> type[StructDataclass]: ... - - -@overload -def struct_dataclass(_cls: None) -> Callable[[type[StructDataclass]], type[StructDataclass]]: ... - - -def struct_dataclass( - _cls: type[StructDataclass] | None = None, -) -> Callable[[type[StructDataclass]], type[StructDataclass]] | type[StructDataclass]: - def inner(_cls: type[StructDataclass]) -> type[StructDataclass]: - new_cls = _cls - - # new_cls should not already be a dataclass - if is_dataclass(new_cls): - return cast(type[StructDataclass], new_cls) - - # Make sure any fields without a default have one - for type_iterator in iterate_types(new_cls): - if not type_iterator.is_pystructtype: - continue - - if not type_iterator.type_meta or type_iterator.type_meta.size == 1: - if type_iterator.is_list: - raise Exception("You said this should be size 1, so this shouldn't be a list") - - # Set a default if it does not yet exist - if not getattr(new_cls, type_iterator.key, None): - default: type | int | float = type_iterator.base_type - if type_iterator.type_meta and type_iterator.type_meta.default: - default = type_iterator.type_meta.default - if isinstance(default, list): - raise Exception("A default for a size 1 should not be a list") - - # Create a new instance of the class - if inspect.isclass(default): - default = field(default_factory=lambda d=default: d()) # type: ignore - else: - default = field(default_factory=lambda d=default: deepcopy(d)) # type: ignore - - setattr(new_cls, type_iterator.key, default) - else: - # This assumes we want multiple items of base_type, so make sure the given base_type is - # properly set to be a list as well - if not type_iterator.is_list: - raise Exception("You want a list, so make it a list you dummy") - - # We have a meta type and the size is > 1 so make the default a field - default = type_iterator.base_type - if type_iterator.type_meta and type_iterator.type_meta.default: - default = type_iterator.type_meta.default - - default_list = [] - if isinstance(default, list): - # TODO: Implement having the entire list be a default rather than needing to set each - # TODO: element as the same base object. - pass - else: - # Create a new instance of the class - if inspect.isclass(default): - default_list = field( - default_factory=lambda d=default, s=type_iterator.type_meta.size: [ # type: ignore - d() for _ in range(s) - ] - ) - else: - default_list = field( - default_factory=lambda d=default, s=type_iterator.type_meta.size: [ # type: ignore - deepcopy(d) for _ in range(s) - ] - ) - - setattr(new_cls, type_iterator.key, default_list) - return cast(type[StructDataclass], dataclass(new_cls)) - - if _cls is None: - return inner - return inner(_cls) - - -def int_to_bool_list(data: int | list[int], byte_length: int) -> list[bool]: - """ - Converts integer or a list of integers into a list of bools representing the bits - - ex. ord("A") = [False, True, False, False, False, False, False, True] - - ex. [ord("A"), ord("B")] = [False, True, False, False, False, False, False, True, - False, True, False, False, False, False, True, False] - - :param data: Integer(s) to be converted - :param byte_length: Number of bytes to extract from integer(s) - :return: List of bools representing each bit in the data - """ - # Convert a single int into a list, so we can assume we're always working with a list here - data = [data] if isinstance(data, int) else data - - # The amount of bits we end up with will be the number of bytes we expect in the int times 8 (8 bits in a byte) - # For example uint8_t would have 1 byte, but uint16_t would have 2 bytes - byte_size = (byte_length * 8) // len(data) - - bit_strs = [] - for val in data: - # Convert the int(s) in to a string of bits (add 2 to account for the `0b` prefix) - tmp_str = format(val, f"#0{byte_size + 2}b") - # Cut off the `0b` prefix of the bit string, and reverse it - bit_strs.append(tmp_str.removeprefix("0b")[::-1]) - # Convert the bit_str to a list of ints - bit_list = map(int, "".join(bit_strs[::-1])) - # Convert the bit list to bools and return - return list(map(bool, bit_list)) - - -class BitsType(StructDataclass): - _raw: Any - _meta: dict - _meta_tuple: tuple - - def __post_init__(self): - super().__post_init__() - - self._meta = {k: v for k, v in zip(*self._meta_tuple, strict=False)} - - def _decode(self, data: list[int]) -> None: - # First call the super function to put the values in to _raw - super()._decode(data) - - # Combine all data in _raw as binary and convert to bools - bin_data = int_to_bool_list(self._raw, self._byte_length) - - for k, v in self._meta.items(): - if isinstance(v, list): - steps = [] - for idx in v: - steps.append(bin_data[idx]) - setattr(self, k, steps) - else: - setattr(self, k, bin_data[v]) - - def _encode(self) -> list[int]: - bin_data = list(itertools.repeat(False, self._byte_length * 8)) - for k, v in self._meta.items(): - if isinstance(v, list): - steps = getattr(self, k) - for idx, bit_idx in enumerate(v): - bin_data[bit_idx] = steps[idx] - else: - bin_data[v] = getattr(self, k) - - if isinstance(self._raw, list): - self._raw = [ - sum(v << i for i, v in enumerate(chunk)) - for chunk in list_chunks(bin_data, (self._byte_length // len(self._raw)) * 8) - ][::-1] - else: - self._raw = sum(v << i for i, v in enumerate(bin_data)) - - # Run the super function to return the data in self._raw() - return super()._encode() - - -def bits(_type: Any, definition: dict[str, int | list[int]]) -> Callable[[type[BitsType]], type[StructDataclass]]: - def inner(_cls: type[BitsType]) -> type[StructDataclass]: - # Create class attributes based on the definition - # TODO: Maybe a sanity check to make sure the definition is the right format, and no overlapping bits, etc - - new_cls = _cls - - new_cls.__annotations__["_raw"] = _type - - new_cls._meta = field(default_factory=dict) - new_cls.__annotations__["_meta"] = dict[str, int] - - # Convert the definition to a named tuple, so it's Immutable - meta_tuple = (tuple(definition.keys()), tuple(definition.values())) - new_cls._meta_tuple = field(default_factory=lambda d=meta_tuple: d) # type: ignore - new_cls.__annotations__["_meta_tuple"] = tuple - - # TODO: Support int, or list of ints as defaults - # TODO: Support dict, and dict of lists, or list of dicts, etc for definition - # TODO: ex. definition = {"a": {"b": 0, "c": [1, 2, 3]}, "d": [4, 5, 6], "e": {"f": 7}} - # TODO: Can't decide if the line above this is a good idea or not - for key, value in definition.items(): - if isinstance(value, list): - setattr( - new_cls, - key, - field(default_factory=lambda v=len(value): [False for _ in range(v)]), # type: ignore # noqa: B008 - ) - new_cls.__annotations__[key] = list[bool] - else: - setattr(new_cls, key, False) - new_cls.__annotations__[key] = bool - - return struct_dataclass(new_cls) - - return inner - - -# XXX: This is how class decorators essentially work +__all__ = [ + "BitsType", + "StructDataclass", + "TypeInfo", + "TypeMeta", + "bits", + "char_t", + "double_t", + "float_t", + "int8_t", + "int16_t", + "int32_t", + "int64_t", + "struct_dataclass", + "uint8_t", + "uint16_t", + "uint32_t", + "uint64_t", +] + +# Note: This is how class decorators essentially work +# # @foo # class gotem(): ... # diff --git a/src/pystructtype/bitstype.py b/src/pystructtype/bitstype.py new file mode 100644 index 0000000..58b6d22 --- /dev/null +++ b/src/pystructtype/bitstype.py @@ -0,0 +1,142 @@ +import itertools +from collections.abc import Callable +from dataclasses import field +from typing import Any + +from pystructtype.structdataclass import StructDataclass, struct_dataclass +from pystructtype.utils import int_to_bool_list, list_chunks + + +class BitsType(StructDataclass): + """ + Class to auto-magically decode/encode struct data into separate variables + for separate bits based on the given definition + """ + + _raw: Any + _meta: dict + _meta_tuple: tuple + + def __post_init__(self) -> None: + super().__post_init__() + + # Convert the _meta_tuple data into a dictionary and put it into _meta + self._meta = {k: v for k, v in zip(*self._meta_tuple, strict=False)} + + def _decode(self, data: list[int]) -> None: + """ + Internal decoding function + + :param data: A list of ints to decode + """ + # First call the super function to put the values in to _raw + super()._decode(data) + + # Combine all data in _raw as binary and convert to bools + bin_data = int_to_bool_list(self._raw, self._byte_length) + + # Apply bits to the defined structure + for k, v in self._meta.items(): + if isinstance(v, list): + steps = [] + for idx in v: + steps.append(bin_data[idx]) + setattr(self, k, steps) + else: + setattr(self, k, bin_data[v]) + + def _encode(self) -> list[int]: + """ + Internal encoding function + + :returns: A list of encoded ints + """ + # Fill a correctly sized variable with all False/0 bits + bin_data = list(itertools.repeat(False, self._byte_length * 8)) + + # Assign the correct values from the defined attributes into bin_data + for k, v in self._meta.items(): + if isinstance(v, list): + steps = getattr(self, k) + for idx, bit_idx in enumerate(v): + bin_data[bit_idx] = steps[idx] + else: + bin_data[v] = getattr(self, k) + + # Convert bin_data back into their correct integer locations + if isinstance(self._raw, list): + self._raw = [ + sum(v << i for i, v in enumerate(chunk)) + for chunk in list_chunks(bin_data, (self._byte_length // len(self._raw)) * 8) + ][::-1] + else: + self._raw = sum(v << i for i, v in enumerate(bin_data)) + + # Run the super function to return the data in self._raw() + return super()._encode() + + +def bits(_type: Any, definition: dict[str, int | list[int]]) -> Callable[[type[BitsType]], type[StructDataclass]]: + """ + Decorator that does a bunch of metaprogramming magic to properly set up the + defined Subclass of StructDataclass for Bits handling + + The definition must be a dict of ints or a list of ints. The int values denote the position of the bits. + + Example: + @bits(uint8_t, {"a": 0, "b": [1, 2, 4], "c": 3}) + class MyBits(BitsType): ... + + For an uint8_t defined as 0b01010101, the resulting class will be: + MyBits(a=1, b=[0, 1, 1], c=0) + + :param _type: The type of data that the bits are stored in (ex. uint8_t, etc.) + :param definition: The bits definition that defines attributes and bit locations + :return: A Callable that performs the metaprogramming magic and returns the modified StructDataclass + """ + + def inner(_cls: type[BitsType]) -> type[StructDataclass]: + """ + The inner function to modify a StructDataclass into a BitsType class + + :param _cls: A Subclass of BitsType + :return: Modified StructDataclass + """ + # Create class attributes based on the definition + # TODO: Maybe a sanity check to make sure the definition is the right format, and no overlapping bits, etc + + new_cls = _cls + + # Set the correct type for the raw data + new_cls.__annotations__["_raw"] = _type + + # Override the annotations for the _meta attribute, and set a default + # TODO: This probably isn't really needed unless we end up changing the int value to bool or something + new_cls._meta = field(default_factory=dict) + new_cls.__annotations__["_meta"] = dict[str, int] + + # Convert the definition to a named tuple, so it's Immutable + meta_tuple = (tuple(definition.keys()), tuple(definition.values())) + new_cls._meta_tuple = field(default_factory=lambda d=meta_tuple: d) # type: ignore + new_cls.__annotations__["_meta_tuple"] = tuple + + # TODO: Support int, or list of ints as defaults + # TODO: Support dict, and dict of lists, or list of dicts, etc for definition + # TODO: ex. definition = {"a": {"b": 0, "c": [1, 2, 3]}, "d": [4, 5, 6], "e": {"f": 7}} + # TODO: Can't decide if the line above this is a good idea or not + # Create the defined attributes, defaults, and annotations in the class + for key, value in definition.items(): + if isinstance(value, list): + setattr( + new_cls, + key, + field(default_factory=lambda v=len(value): [False for _ in range(v)]), # type: ignore # noqa: B008 + ) + new_cls.__annotations__[key] = list[bool] + else: + setattr(new_cls, key, False) + new_cls.__annotations__[key] = bool + + return struct_dataclass(new_cls) + + return inner diff --git a/src/pystructtype/structdataclass.py b/src/pystructtype/structdataclass.py new file mode 100644 index 0000000..196ec4f --- /dev/null +++ b/src/pystructtype/structdataclass.py @@ -0,0 +1,362 @@ +import inspect +import re +import struct +from collections.abc import Callable +from copy import deepcopy +from dataclasses import dataclass, field, is_dataclass +from typing import cast, overload + +from pystructtype.structtypes import iterate_types + + +@dataclass +class StructState: + """ + Contains necessary struct information to correctly + decode and encode the data in a StructDataclass + """ + + name: str + struct_fmt: str + size: int + + +class StructDataclass: + """ + Class that will auto-magically decode and encode data for the defined + subclass. + """ + + def __post_init__(self) -> None: + self._state: list[StructState] = [] + + # Grab Struct Format + self.struct_fmt = "" + for type_iterator in iterate_types(self.__class__): + if type_iterator.type_info: + self._state.append( + StructState( + type_iterator.key, + type_iterator.type_info.format, + type_iterator.size, + ) + ) + self.struct_fmt += ( + f"{type_iterator.size if type_iterator.size > 1 else ''}{type_iterator.type_info.format}" + ) + elif inspect.isclass(type_iterator.base_type) and issubclass(type_iterator.base_type, StructDataclass): + attr = getattr(self, type_iterator.key) + if type_iterator.is_list: + fmt = attr[0].struct_fmt + else: + fmt = attr.struct_fmt + self._state.append(StructState(type_iterator.key, fmt, type_iterator.size)) + self.struct_fmt += fmt * type_iterator.size + else: + # We have no TypeInfo object, and we're not a StructDataclass + # This means we're a regularly defined class variable, and we + # Don't have to do anything about this. + pass + self._simplify_format() + self._byte_length = struct.calcsize("=" + self.struct_fmt) + # print(f"{self.__class__.__name__}: {self._byte_length} : {self.struct_fmt}") + + def _simplify_format(self) -> None: + """ + Simplify the struct format that has been defined for this class. + + Essentially we turn things like `ccbbbbh` into `2c4bh` + """ + # Expand any already condensed sections + # This can happen if we have nested StructDataclasses + expanded_format = "" + items = re.findall(r"([a-zA-Z]|\d+)", self.struct_fmt) + items_len = len(items) + idx = 0 + while idx < items_len: + if "0" <= (item := items[idx]) <= "9": + idx += 1 + expanded_format += items[idx] * int(item) + else: + expanded_format += item + idx += 1 + + # Simplify the format by turning multiple consecutive letters into a number + letter combo + simplified_format = "" + for group in (x[0] for x in re.findall(r"(([a-zA-Z])\2*)", expanded_format)): + simplified_format += f"{group_len if (group_len := len(group)) > 1 else ''}{group[0]}" + + self.struct_fmt = simplified_format + + def size(self) -> int: + """ + The size of this struct is defined as the sum of the sizes of all attributes + + :return: Combined size of the struct + """ + return sum(state.size for state in self._state) + + @staticmethod + def _endian(little_endian: bool) -> str: + """ + Return "<" or ">" depending on endianness, to pass to struct decode/encode + + :param little_endian: True if we expect little_endian, else False + :return: "<" if little_endian else ">" + """ + return "<" if little_endian else ">" + + @staticmethod + def _to_bytes(data: list[int] | bytes) -> bytes: + """ + Convert a list of ints into bytes + + :param data: a list of ints or a bytes object + :return: a bytes object + """ + if isinstance(data, bytes): + return data + return bytes(data) + + @staticmethod + def _to_list(data: list[int] | bytes) -> list[int]: + """ + Convert a bytes object into a list of ints + + :param data: a list of ints or a bytes object + :return: a list of ints + """ + if isinstance(data, bytes): + return list(data) + return data + + def _decode(self, data: list[int]) -> None: + """ + Internal decoding function for the StructDataclass. + + Extend this function if you wish to add extra processing to your StructDataclass decoding processing + + :param data: A list of ints to decode into the StructDataclass + """ + idx = 0 + for state in self._state: + attr = getattr(self, state.name) + + if isinstance(attr, list) and isinstance(attr[0], StructDataclass): + # If the current attribute is a list, and contains subclasses of StructDataclass + # Call _decode on the required subset of bytes for each list item + list_idx = 0 + sub_struct_byte_length = attr[0].size() + while list_idx < state.size: + instance: StructDataclass = attr[list_idx] + instance._decode(data[idx : idx + sub_struct_byte_length]) + list_idx += 1 + idx += sub_struct_byte_length + elif isinstance(attr, StructDataclass): + # If the current attribute is not a list, and is a subclass of StructDataclass + # Call _decode on the required subset of bytes for the item + if state.size != 1: + raise Exception(f"Attribute {state.name} is not defined as a list but has a size > 0") + + sub_struct_byte_length = attr.size() + attr._decode(data[idx : idx + sub_struct_byte_length]) + idx += sub_struct_byte_length + elif state.size == 1: + # The current attribute is a base type of size 1 + setattr(self, state.name, data[idx]) + idx += 1 + else: + # The current attribute is a list of base types + list_idx = 0 + while list_idx < state.size: + getattr(self, state.name)[list_idx] = data[idx] + list_idx += 1 + idx += 1 + + def decode(self, data: list[int] | bytes, little_endian=False) -> None: + """ + Decode the given data into this subclass of StructDataclass + + :param data: list of ints or a bytes object + :param little_endian: True if decoding little_endian formatted data, else False + """ + data = self._to_bytes(data) + + # Decode + self._decode(list(struct.unpack(self._endian(little_endian) + self.struct_fmt, data))) + + def _encode(self) -> list[int]: + """ + Internal encoding function for the StructDataclass. + + Extend this function if you wish to add extra processing to your StructDataclass encoding processing + + :return: list of encoded int data + """ + result: list[int] = [] + + for state in self._state: + attr = getattr(self, state.name) + + if isinstance(attr, list) and isinstance(attr[0], StructDataclass): + # Attribute is a list of StructDataclass subclasses. + # Simply call _encode on each item in the list + item: StructDataclass + for item in attr: + result.extend(item._encode()) + elif isinstance(attr, StructDataclass): + # Attribute is a StructDataclass subclass + # Call _encode on it + if state.size != 1: + raise Exception(f"Attribute {state.name} is defined as a list but has a size == 1") + result.extend(attr._encode()) + elif state.size == 1: + # Attribute is a single base type + # Append it to the result + result.append(getattr(self, state.name)) + else: + # Attribute is a list of base types + # Extend it to the result + result.extend(getattr(self, state.name)) + return result + + def encode(self, little_endian=False) -> bytes: + """ + Encode the data from this subclass of StructDataclass into bytes + + :param little_endian: True if encoding little_endian formatted data, else False + :return: encoded bytes + """ + result = self._encode() + return struct.pack(self._endian(little_endian) + self.struct_fmt, *result) + + +@overload +def struct_dataclass(_cls: type[StructDataclass]) -> type[StructDataclass]: + """ + Overload for using a bare decorator + + @struct_dataclass + class foo(StructDataclass): ... + + Equivalent to: struct_dataclass(foo) + + :param _cls: Subtype of StructDataclass + :return: Modified Subtype of StructDataclass + """ + pass + + +@overload +def struct_dataclass(_cls: None) -> Callable[[type[StructDataclass]], type[StructDataclass]]: + """ + Overload for using called decorator + + @struct_dataclass() + class foo(StructDataclass): ... + + Equivalent to: struct_dataclass()(foo) + + :param _cls: None + :return: Callable that takes in a Subtype of StructDataclass and returns a modified Subtype + """ + pass + + +def struct_dataclass( + _cls: type[StructDataclass] | None = None, +) -> Callable[[type[StructDataclass]], type[StructDataclass]] | type[StructDataclass]: + """ + Decorator that does a bunch of metaprogramming magic to properly set up + the defined Subclass of a StructDataclass + + :param _cls: A Subclass of StructDataclass or None + :return: A Modified Subclass of a StructDataclass or a Callable that performs the same actions + """ + + def inner(_cls: type[StructDataclass]) -> type[StructDataclass]: + """ + The inner function for `struct_dataclass` that actually does all the work + + :param _cls: A Subclass of StructDataclass + :return: A Modified Subclass of a StructDataclass + """ + new_cls = _cls + + # new_cls should not already be a dataclass, + # but it will be a subtype of Dataclass by the end of this function + if is_dataclass(new_cls): + # Just try to cast it again, and return + return cast(type[StructDataclass], new_cls) + + # Make sure any fields without a default have one + # This prevents Dataclass from being mad that we might have attributes defined with + # defaults interwoven between ones that don't + for type_iterator in iterate_types(new_cls): + # If the current type is just a base type, then we can essentially ignore it + # These are typically used for extra processing and not included in the decode/encode + if not type_iterator.is_pystructtype: + continue + + if not type_iterator.type_meta or type_iterator.type_meta.size == 1: + # This type either has no metadata, or is defined as having a size of 1 and is + # therefore not a list + if type_iterator.is_list: + raise Exception(f"Attribute {type_iterator.key} is defined as a list type but has size set to 1") + + # Set a default if it does not yet exist + if not getattr(new_cls, type_iterator.key, None): + default: type | int | float = type_iterator.base_type + if type_iterator.type_meta and type_iterator.type_meta.default: + default = type_iterator.type_meta.default + if isinstance(default, list): + raise Exception(f"default value for {type_iterator.key} attribute can not be a list") + + # Create a new instance of the class, or value + if inspect.isclass(default): + default = field(default_factory=lambda d=default: d()) # type: ignore + else: + default = field(default_factory=lambda d=default: deepcopy(d)) # type: ignore + + setattr(new_cls, type_iterator.key, default) + else: + # This assumes we want multiple items of base_type, so make sure the given base_type is + # properly set to be a list as well + if not type_iterator.is_list: + raise Exception(f"Attribute {type_iterator.key} is not a list type but has a size > 1") + + # We have a meta type and the size is > 1 so make the default a field + default = type_iterator.base_type + if type_iterator.type_meta and type_iterator.type_meta.default: + default = type_iterator.type_meta.default + + default_list = [] + if isinstance(default, list): + # TODO: Implement having the entire list be a default rather than needing to set each + # TODO: element as the same base object. + pass + else: + # Create a new instance of the class or value + if inspect.isclass(default): + default_list = field( + default_factory=lambda d=default, s=type_iterator.type_meta.size: [ # type: ignore + d() for _ in range(s) + ] + ) + else: + default_list = field( + default_factory=lambda d=default, s=type_iterator.type_meta.size: [ # type: ignore + deepcopy(d) for _ in range(s) + ] + ) + + setattr(new_cls, type_iterator.key, default_list) + return cast(type[StructDataclass], dataclass(new_cls)) + + # If we use the decorator with empty parens, we simply return the inner callable + if _cls is None: + return inner + + # If we use the decorator with no parens, we return the result of passing the _cls + # to the inner callable + return inner(_cls) diff --git a/src/pystructtype/structtypes.py b/src/pystructtype/structtypes.py new file mode 100644 index 0000000..5c9b419 --- /dev/null +++ b/src/pystructtype/structtypes.py @@ -0,0 +1,157 @@ +import inspect +from collections.abc import Generator +from dataclasses import dataclass +from typing import Annotated, Any, TypeVar, get_args, get_origin, get_type_hints + +from pystructtype import structdataclass + +T = TypeVar("T", int, float, default=int) +"""Generic Data Type for StructDataclass Contents""" + + +@dataclass(frozen=True) +class TypeMeta[T]: + """ + Class used to define Annotated Type Metadata for + size and default values + """ + + size: int = 1 + default: T | None = None + + +@dataclass(frozen=True) +class TypeInfo: + """ + Class used to define Annotated Type Metadata + for format and byte size + """ + + format: str + byte_size: int + + +# TODO: Support proper "c-string" types + +# Fixed Size Types +char_t = Annotated[int, TypeInfo("c", 1)] +"""1 Byte char Type""" +int8_t = Annotated[int, TypeInfo("b", 1)] +"""1 Byte Signed int Type""" +uint8_t = Annotated[int, TypeInfo("B", 1)] +"""1 Byte Unsigned int Type""" +int16_t = Annotated[int, TypeInfo("h", 2)] +"""2 Byte Signed int Type""" +uint16_t = Annotated[int, TypeInfo("H", 2)] +"""2 Byte Unsigned int Type""" +int32_t = Annotated[int, TypeInfo("i", 4)] +"""4 Byte Signed int Type""" +uint32_t = Annotated[int, TypeInfo("I", 4)] +"""4 Byte Unsigned int Type""" +int64_t = Annotated[int, TypeInfo("q", 8)] +"""8 Byte Signed int Type""" +uint64_t = Annotated[int, TypeInfo("Q", 8)] +"""8 Byte Unsigned int Type""" + +# TODO: Make a special Bool class to auto-convert from int to bool + +# Named Types +float_t = Annotated[float, TypeInfo("f", 4)] +"""4 Byte float Type""" +double_t = Annotated[float, TypeInfo("d", 8)] +"""8 Byte double Type""" + + +@dataclass +class TypeIterator: + """ + Contains all relevant type information for + an object in a StructDataclass. + + Used as a container when iterating through StructDataclass attributes + """ + + key: str + base_type: type + type_info: TypeInfo | None + type_meta: TypeMeta | None + is_list: bool + is_pystructtype: bool + + @property + def size(self) -> int: + """ + Return the size of the type. If this is not a list, this will default to + 1, else this will return the size defined in the `type_meta` object + if it exists. + + :return: integer containing the size of the type + """ + return getattr(self.type_meta, "size", 1) + + +def iterate_types(cls: type) -> Generator[TypeIterator]: + """ + Iterate through the given StructDataclass attributes type hints and yield + a TypeIterator for each one. + + :param cls: A StructDataclass class object (not an instantiated object) + :return: Yield a TypeIterator object + """ + for key, hint in get_type_hints(cls, include_extras=True).items(): + # Grab the base type from a possibly annotated type hint + base_type = type_from_annotation(hint) + + # Determine if the type is a list + # ex. list[bool] (yes) vs bool (no) + is_list = issubclass(origin, list) if (origin := get_origin(base_type)) else False + + # Grab the first args value and look for any TypeMeta objects within + type_args = get_args(hint) + type_meta = next((x for x in type_args if isinstance(x, TypeMeta)), None) + + # type_args has the possibility of being nested within more tuples + # drill down the type_args until we hit empty, then we know we're at the bottom + # which is where type_info will exist + if type_args and len(type_args) > 1: + while args := get_args(type_args[0]): + type_args = args + + # Find the TypeInfo object on the lowest rung of the type_args + type_info = next((x for x in type_args if isinstance(x, TypeInfo)), None) + + # At this point we may have possibly drilled down into `type_args` to find the true base type + if type_args: + base_type = type_from_annotation(type_args[0]) + + # Determine if we are a subclass of a pystructtype: + # A pystructtype will be a type with a type_info object in the Annotation, + # or a subtype of StructDataclass + is_pystructtype = type_info is not None or ( + inspect.isclass(base_type) and issubclass(base_type, structdataclass.StructDataclass) + ) + + yield TypeIterator(key, base_type, type_info, type_meta, is_list, is_pystructtype) + + +def type_from_annotation(_type: type) -> type: + """ + Find the base type from an Annotated type, + or return it unchanged if not Annotated + + :param _type: Type to check + :return: Base type if Annotated, or the original passed in type otherwise + """ + # If we have an origin for the given type, and it's Annotated + if (origin := get_origin(_type)) and origin is Annotated: + # Keep running `get_args` on the first element of whatever + # `get_args` returns, until we get nothing back + arg = _type + t: Any = _type + while t := get_args(t): + arg = t[0] + + # This will be the base type + return arg + # No origin, or the origin is not Annotated, just return the given type + return _type diff --git a/src/pystructtype/utils.py b/src/pystructtype/utils.py new file mode 100644 index 0000000..db751b9 --- /dev/null +++ b/src/pystructtype/utils.py @@ -0,0 +1,43 @@ +from collections.abc import Generator + + +def list_chunks(_list: list, n: int) -> Generator[list]: + """ + Yield successive n-sized chunks from a list. + :param _list: List to chunk out + :param n: Size of chunks + :return: Generator of n-sized chunks of _list + """ + yield from (_list[i : i + n] for i in range(0, len(_list), n)) + + +def int_to_bool_list(data: int | list[int], byte_length: int) -> list[bool]: + """ + Converts an integer or a list of integers into a list of bools representing the bits + + ex. ord("A") or 0b01000001 = [False, True, False, False, False, False, False, True] + + ex. [ord("A"), ord("B")] = [False, True, False, False, False, False, False, True, + False, True, False, False, False, False, True, False] + + :param data: Integer(s) to be converted + :param byte_length: Number of bytes to extract from integer(s) + :return: List of bools representing each bit in the data + """ + # Convert a single int into a list, so we can assume we're always working with a list here + data = [data] if isinstance(data, int) else data + + # The amount of bits we end up with will be the number of bytes we expect in the int times 8 (8 bits in a byte) + # For example uint8_t would have 1 byte, but uint16_t would have 2 bytes + byte_size = (byte_length * 8) // len(data) + + bit_strs = [] + for val in data: + # Convert the int(s) in to a string of bits (add 2 to account for the `0b` prefix) + tmp_str = format(val, f"#0{byte_size + 2}b") + # Cut off the `0b` prefix of the bit string, and reverse it + bit_strs.append(tmp_str.removeprefix("0b")[::-1]) + # Convert the bit_str to a list of ints representing single bits + bit_list = map(int, "".join(bit_strs[::-1])) + # Convert the bit list to bools and return + return list(map(bool, bit_list)) diff --git a/test/examples.py b/test/examples.py index eb53583..98d0cfb 100644 --- a/test/examples.py +++ b/test/examples.py @@ -3,16 +3,8 @@ from enum import IntEnum from typing import Annotated -from pystructtype import ( - BitsType, - StructDataclass, - TypeMeta, - bits, - list_chunks, - struct_dataclass, - uint8_t, - uint16_t, -) +from pystructtype import BitsType, StructDataclass, TypeMeta, bits, struct_dataclass, uint8_t, uint16_t +from pystructtype.utils import list_chunks TEST_CONFIG_DATA = [ # masterVersion diff --git a/test/test_ctypes.py b/test/test_ctypes.py index 0590716..8d78d14 100644 --- a/test/test_ctypes.py +++ b/test/test_ctypes.py @@ -1,13 +1,6 @@ from typing import Annotated -from pystructtype import ( - BitsType, - StructDataclass, - TypeMeta, - bits, - struct_dataclass, - uint8_t, -) +from pystructtype import BitsType, StructDataclass, TypeMeta, bits, struct_dataclass, uint8_t from .examples import TEST_CONFIG_DATA, SMXConfigType # type: ignore