import dataclasses import textwrap import typing as t import construct as cs from .generic_wrapper import ( Adapter, BuildTypes, Construct, Context, ParsedType, PathType, ) if t.TYPE_CHECKING: class _TContainerBase(cs.Container[t.Any]): ... else: class _TContainerBase(cs.Container): pass class TContainerBase(_TContainerBase): """ Base class for a Container of a TStruct and a TBitStruct. Note: this always has to be mixed with "dataclasses.dataclass". """ def __getattribute__(self, name: str) -> t.Any: # if accessing via an field via dot access, return the object from the dict if name in self: return self[name] else: return super().__getattribute__(name) def __post_init__(self) -> None: # 1. append fields with init=False to the dict of the cs.Container # 2. fix the order of the OrderedDict for field in dataclasses.fields(self): value = getattr(self, field.name) if field.init is False: self[field.name] = value else: self.move_to_end(field.name) def sfield( subcon: Construct[ParsedType, t.Any], doc: t.Optional[str] = None, parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, ) -> ParsedType: """ Create a dataclass field for a "TStruct" and "TBitStruct" from a subcon. """ # Rename subcon, if doc or parsed are available if (doc is not None) or (parsed is not None): if doc is not None: doc = textwrap.dedent(doc).strip("\n") subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) if subcon.flagbuildnone is True: # if subcon builds from "None", set default to "None" field = dataclasses.field( default=None, init=False, metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, ) else: field = dataclasses.field(metadata={"subcon": subcon}) return field # type: ignore TStructField = sfield # also support legacy name ContainerType = t.TypeVar("ContainerType", bound=TContainerBase) class _TStruct(Adapter[t.Any, t.Any, ContainerType, BuildTypes]): """ Base class for a typed struct, based on standard dataclasses. """ def __init__( self, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False, ) -> None: if not issubclass(container_type, TContainerBase): raise TypeError( "'{}' has to be a '{}'".format( repr(container_type), repr(TContainerBase) ) ) if not dataclasses.is_dataclass(container_type): raise TypeError( "'{}' has to be a 'dataclasses.dataclass'".format(repr(container_type)) ) self.container_type = container_type self.swapped = swapped self.add_offsets = add_offsets # get all fields from the dataclass fields = dataclasses.fields(self.container_type) if self.swapped: fields = tuple(reversed(fields)) # extract the construct formats from the struct_type subcon_fields = {} for field in fields: if add_offsets: subcon_fields[f"@<{field.name}"] = cs.Tell subcon_fields[field.name] = field.metadata["subcon"] if add_offsets: subcon_fields[f"@>{field.name}"] = cs.Tell # init adatper super(_TStruct, self).__init__(self._create_subcon(subcon_fields)) # type: ignore def _create_subcon( self, subcon_fields: t.Dict[str, t.Any] ) -> Construct[t.Any, t.Any]: raise NotImplementedError def __getattr__(self, name: str) -> t.Any: return getattr(self.subcon, name) def _decode( self, obj: "cs.Container[t.Any]", context: Context, path: PathType ) -> ContainerType: # get all fields from the dataclass fields = dataclasses.fields(self.container_type) # extract all fields from the container, that are used for create the dataclass object dc_init = {} for field in fields: if field.init: value = getattr(obj, field.name) dc_init[field.name] = value # create object of dataclass dc = self.container_type(**dc_init) # extract all other values from the container, an pass it to the dataclass for field in fields: if not field.init: value = getattr(obj, field.name) setattr(dc, field.name, value) # also transfer values that are not part of the dataclass dc.update(obj) return dc def _encode( self, obj: BuildTypes, context: Context, path: PathType ) -> t.Dict[str, t.Any]: if isinstance(obj, self.container_type): # get all fields from the dataclass fields = dataclasses.fields(self.container_type) # extract all fields from the container, that are used for create the dataclass object ret_dict = {} for field in fields: value = getattr(obj, field.name) ret_dict[field.name] = value return ret_dict raise TypeError( "'{}' has to be of type {}".format(repr(obj), repr(self.container_type)) ) class TStruct(_TStruct[ContainerType, ContainerType]): """ Typed struct, based on standard dataclasses. """ subcon: "cs.Struct[t.Any, t.Any]" if t.TYPE_CHECKING: def __new__( cls, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False, ) -> "TStruct[ContainerType]": ... def _create_subcon( self, subcon_fields: t.Dict[str, t.Any] ) -> Construct[t.Any, t.Any]: return cs.Struct(**subcon_fields) class TBitStruct(_TStruct[ContainerType, ContainerType]): """ Typed bit struct, based on standard dataclasses. """ if t.TYPE_CHECKING: def __new__( cls, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False, ) -> "TBitStruct[ContainerType]": ... def _create_subcon( self, subcon_fields: t.Dict[str, t.Any] ) -> Construct[t.Any, t.Any]: return cs.BitStruct(**subcon_fields)