diff --git a/construct_typed/__init__ old.py b/construct_typed/__init__ old.py new file mode 100644 index 0000000..24e0e36 --- /dev/null +++ b/construct_typed/__init__ old.py @@ -0,0 +1,288 @@ +import enum +import typing as t +import textwrap +import construct as cs +import dataclasses +import .generic_construct as gcs + +ParsedType = t.TypeVar("ParsedType") +BuildTypes = t.TypeVar("BuildTypes") +SubconParsedType = t.TypeVar("SubconParsedType") +SubconBuildTypes = t.TypeVar("SubconBuildTypes") +EnumType = t.TypeVar("EnumType", bound=enum.IntEnum) +DataclassType = t.TypeVar("DataclassType") +ListType = t.TypeVar("ListType") + + +# if t.TYPE_CHECKING: +# # while type checking, the original classes are generics, because they are defined in the stubs. +# from construct import Construct, Adapter +# from construct import ListContainer as List +# else: +# # at runtime, the original classes are no generics, so whe have to make new classes with generics support +# class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct): +# pass + +# class Adapter( +# t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes], +# cs.Adapter, +# ): +# pass + +# class List(t.Generic[ListType], cs.ListContainer): +# pass + + + +# =============================================================================== +# mappings +# =============================================================================== +class TEnum(gcs.Adapter[int, int, EnumType, t.Union[int, str, EnumType]]): + def __new__( + cls, subcon: gcs.Construct[int, int], enum_type: t.Type[EnumType] + ) -> "TEnum[EnumType]": + return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore + + def __init__(self, subcon: gcs.Construct[int, int], enum_type: t.Type[EnumType]): + if not issubclass(enum_type, enum.IntEnum): + raise TypeError( + 'The class "{}" is not an "enum.IntEnum"'.format( + enum_type.__name__ + ) + ) + + @classmethod + def _missing_( + cls: t.Type[EnumType], value: t.Union[int, EnumType] + ) -> t.Optional[EnumType]: + if isinstance(value, int): + return cls._create_pseudo_member_(value) + return None # will raise the ValueError in Enum.__new__ + + @classmethod + def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType: + pseudo_member = cls._value2member_map_.get(value, None) + if pseudo_member is None: + new_member = int.__new__(cls, value) + # I expect a name attribute to hold a string, hence str(value) + # However, new_member._name_ = value works, too + new_member._name_ = str(value) + new_member._value_ = value + pseudo_member = cls._value2member_map_.setdefault(value, new_member) + return pseudo_member + + # Monkey-patch the enum type with __missing__ method. So if a enum value + # not found in the enum a new pseudo member is created. + # The idea is taken from: https://stackoverflow.com/a/57179436 + enum_type._missing_ = _missing_ + enum_type._create_pseudo_member_ = _create_pseudo_member_ + + # save enum type + self.enum_type = enum_type + + # init adatper + super(TEnum, self).__init__(subcon) # type: ignore + + def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType: + return self.enum_type(obj) + + def _encode( + self, + obj: t.Union[int, str, EnumType], + context: "cs.Context", + path: "cs.PathType", + ) -> int: + try: + if isinstance(obj, str): + return int(self.enum_type[obj]) + else: + return int(self.enum_type(obj)) + except: + raise cs.MappingError( + "building failed, no mapping for %r" % (obj,), path=path + ) + + +# =============================================================================== +# structures and sequences +# =============================================================================== +def StructField( + subcon: gcs.Construct[ParsedType, BuildTypes], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, "cs.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) + subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) + + if subcon.flagbuildnone is True: + # some subcons have a predefined default value. all other have "None" + default: t.Any = None + if isinstance(subcon, (cs.Const, cs.Default)): + if callable(subcon.value): + raise ValueError("lamda as default is not supported") + default = subcon.value + + # if subcon builds from "None", set default to "None" + field = dataclasses.field( + default=default, + init=False, + metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, + ) + else: + field = dataclasses.field(metadata={"subcon": subcon}) + + return field # type: ignore + + +class _TStruct( + gcs.Adapter["cs.Container[t.Any]", t.Dict[str, t.Any], DataclassType, DataclassType] +): + """ + Base class for a typed struct, based on standard dataclasses. + """ + + def __new__( + cls, dataclass_type: t.Type[DataclassType], swapped: bool = False + ) -> "_TStruct[DataclassType]": + return super(_TStruct, cls).__new__(cls, dataclass_type, swapped) # type: ignore + + def __init__( + self, dataclass_type: t.Type[DataclassType], swapped: bool = False + ) -> None: + if not dataclasses.is_dataclass(dataclass_type): + raise TypeError( + 'The class "{}" is not a "dataclasses.dataclass"'.format( + dataclass_type.__name__ + ) + ) + self.dataclass_type = dataclass_type + self.swapped = swapped + + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_type) + if self.swapped: + fields = tuple(reversed(fields)) + + # extract the construct formats from the struct_type + subcon_fields = {} + for field in fields: + subcon_fields[field.name] = field.metadata["subcon"] + + # init adatper + super(_TStruct, self).__init__(self._create_subcon(subcon_fields)) # type: ignore + + def _create_subcon( + self, subcon_fields: t.Dict[str, t.Any] + ) -> gcs.Construct[t.Any, t.Any]: + raise NotImplementedError + + def _decode( + self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType" + ) -> DataclassType: + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_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.dataclass_type(**dc_init) # type: ignore + + # 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) + + return dc + + def _encode( + self, obj: DataclassType, context: "cs.Context", path: "cs.PathType" + ) -> t.Dict[str, t.Any]: + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_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 + + +class TStruct(_TStruct[DataclassType]): + """ + Typed struct, based on standard dataclasses. + """ + + def _create_subcon( + self, subcon_fields: t.Dict[str, t.Any] + ) -> gcs.Construct[t.Any, t.Any]: + return cs.Struct(**subcon_fields) + + +class TBitStruct(_TStruct[DataclassType]): + """ + Typed bit struct, based on standard dataclasses. + """ + + def _create_subcon( + self, subcon_fields: t.Dict[str, t.Any] + ) -> gcs.Construct[t.Any, t.Any]: + return cs.BitStruct(**subcon_fields) + + + +#=============================================================================== +# conditional +#=============================================================================== +def UnionField( + subcon: gcs.Construct[ParsedType, BuildTypes], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None, +) -> ParsedType: + """ + Create a dataclass field for a "TUnion" 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) + subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) + + if subcon.flagbuildnone is True: + # some subcons have a predefined default value. all other have "None" + default: t.Any = None + if isinstance(subcon, (cs.Const, cs.Default)): + if callable(subcon.value): + raise ValueError("lamda as default is not supported") + default = subcon.value + + # if subcon builds from "None", set default to "None" + field = dataclasses.field( + default=default, + init=False, + metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, + ) + else: + field = dataclasses.field(metadata={"subcon": subcon}) + + return field # type: ignore + + + +# TODO: TypedUnion +# TODO: TypedLazyStruct +# TODO: TypedSequence: Based on typing.namedtuple +# TODO: FocusedSeq: Based on typing.namedtuple \ No newline at end of file diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index e7d171c..5b894ab 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -1,235 +1,14 @@ -from enum import IntEnum -import typing as t -import textwrap -import construct as cs -import dataclasses +from .tarray import * +from .tenum import * +from .tstruct import * +from .tunion import * -ParsedType = t.TypeVar("ParsedType") -BuildTypes = t.TypeVar("BuildTypes") -SubconParsedType = t.TypeVar("SubconParsedType") -SubconBuildTypes = t.TypeVar("SubconBuildTypes") -EnumType = t.TypeVar("EnumType", bound=IntEnum) -DataclassType = t.TypeVar("DataclassType") - - -if t.TYPE_CHECKING: - # while type checking, the original classes are generics, because they are defined in the stubs. - from construct import Construct, Adapter -else: - # at runtime, the original classes are no generics, so whe have to make new classes with generics support - class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct): - pass - - class Adapter( - t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes], - cs.Adapter, - ): - pass - - -# =============================================================================== -# mappings -# =============================================================================== -class TEnum(Adapter[int, int, EnumType, t.Union[int, str, EnumType]]): - def __new__( - cls, subcon: Construct[int, int], enum_type: t.Type[EnumType] - ) -> "TEnum[EnumType]": - return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore - - def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): - @classmethod - def _missing_( - cls: t.Type[EnumType], value: t.Union[int, EnumType] - ) -> t.Optional[EnumType]: - if isinstance(value, int): - return cls._create_pseudo_member_(value) - return None # will raise the ValueError in Enum.__new__ - - @classmethod - def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType: - pseudo_member = cls._value2member_map_.get(value, None) - if pseudo_member is None: - new_member = int.__new__(cls, value) - # I expect a name attribute to hold a string, hence str(value) - # However, new_member._name_ = value works, too - new_member._name_ = str(value) - new_member._value_ = value - pseudo_member = cls._value2member_map_.setdefault(value, new_member) - return pseudo_member - - # Monkey-patch the enum type with __missing__ method. So if a enum value - # not found in the enum a new pseudo member is created. - # The idea is taken from: https://stackoverflow.com/a/57179436 - enum_type._missing_ = _missing_ - enum_type._create_pseudo_member_ = _create_pseudo_member_ - - # save enum type - self.enum_type = enum_type - - # init adatper - super(TEnum, self).__init__(subcon) # type: ignore - - def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType: - return self.enum_type(obj) - - def _encode( - self, - obj: t.Union[int, str, EnumType], - context: "cs.Context", - path: "cs.PathType", - ) -> int: - try: - if isinstance(obj, str): - return int(self.enum_type[obj]) - else: - return int(self.enum_type(obj)) - except: - raise cs.MappingError( - "building failed, no mapping for %r" % (obj,), path=path - ) - - -# =============================================================================== -# structures and sequences -# =============================================================================== -def TSubcon( - subcon: Construct[ParsedType, BuildTypes], - doc: t.Optional[str] = None, - parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None, -) -> ParsedType: - """ - Create a dataclass field 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) - subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) - - if subcon.flagbuildnone is True: - # some subcons have a predefined default value. all other have "None" - default: t.Any = None - if isinstance(subcon, (cs.Const, cs.Default)): - if callable(subcon.value): - raise ValueError("lamda as default is not supported") - default = subcon.value - - # if subcon builds from "None", set default to "None" - field = dataclasses.field( - default=default, - init=False, - metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, - ) - else: - field = dataclasses.field(metadata={"subcon": subcon}) - - return field # type: ignore - - -class _TStruct( - Adapter["cs.Container[t.Any]", t.Dict[str, t.Any], DataclassType, DataclassType] -): - """ - Base class for a typed struct, based on standard dataclasses. - """ - - def __new__( - cls, dataclass_type: t.Type[DataclassType], swapped: bool = False - ) -> "_TStruct[DataclassType]": - return super(_TStruct, cls).__new__(cls, dataclass_type, swapped) # type: ignore - - def __init__( - self, dataclass_type: t.Type[DataclassType], swapped: bool = False - ) -> None: - if not dataclasses.is_dataclass(dataclass_type): - raise TypeError( - 'the "dataclass_type" has to be a dataclass but is "{}"'.format( - type(dataclass_type).__name__ - ) - ) - self.dataclass_type = dataclass_type - self.swapped = swapped - - # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_type) - if self.swapped: - fields = tuple(reversed(fields)) - - # extract the construct formats from the struct_type - subcon_fields = {} - for field in fields: - subcon_fields[field.name] = field.metadata["subcon"] - - # 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 _decode( - self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType" - ) -> DataclassType: - # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_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.dataclass_type(**dc_init) # type: ignore - - # 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) - - return dc - - def _encode( - self, obj: DataclassType, context: "cs.Context", path: "cs.PathType" - ) -> t.Dict[str, t.Any]: - # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_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 - - -class TStruct(_TStruct[DataclassType]): - """ - Typed struct, based on standard dataclasses. - """ - - def _create_subcon( - self, subcon_fields: t.Dict[str, t.Any] - ) -> Construct[t.Any, t.Any]: - return cs.Struct(**subcon_fields) - - -class TBitStruct(_TStruct[DataclassType]): - """ - Typed bit struct, based on standard dataclasses. - """ - - def _create_subcon( - self, subcon_fields: t.Dict[str, t.Any] - ) -> Construct[t.Any, t.Any]: - return cs.BitStruct(**subcon_fields) - - -# TODO: TypedUnion -# TODO: TypedLazyStruct -# TODO: TypedSequence: Based on typing.namedtuple -# TODO: FocusedSeq: Based on typing.namedtuple \ No newline at end of file +__all__ = [ + "StructField", + "TStruct", + "TBitStruct", + "TEnum", + "UnionField", + "TUnion", + "TArray" +] diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generic_wrapper.py new file mode 100644 index 0000000..a773cd4 --- /dev/null +++ b/construct_typed/generic_wrapper.py @@ -0,0 +1,43 @@ +import enum +import typing as t +import textwrap +import construct as cs +import dataclasses + +ParsedType = t.TypeVar("ParsedType") +BuildTypes = t.TypeVar("BuildTypes") +SubconParsedType = t.TypeVar("SubconParsedType") +SubconBuildTypes = t.TypeVar("SubconBuildTypes") +DataclassType = t.TypeVar("DataclassType") +ListType = t.TypeVar("ListType") +ValueType = t.TypeVar("ValueType") +EnumType = t.TypeVar("EnumType", bound=enum.IntEnum) + + +if t.TYPE_CHECKING: + # while type checking, the original classes are generics, because they are defined in the stubs. + from construct import Construct as Construct + from construct import Adapter as Adapter + from construct import ListContainer as ListContainer + from construct import Context as Context + from construct import ConstantOrContextLambda as ConstantOrContextLambda + from construct import PathType as PathType +else: + # at runtime, the original classes are no generics, so whe have to make new classes with generics support + class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct): + pass + + class Adapter( + t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes], + cs.Adapter, + ): + pass + + class ListContainer(t.Generic[ListType], cs.ListContainer): + pass + + class Context: + pass + + ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]] + PathType = str diff --git a/construct_typed/py.typed b/construct_typed/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/construct_typed/tarray.py b/construct_typed/tarray.py new file mode 100644 index 0000000..ced182a --- /dev/null +++ b/construct_typed/tarray.py @@ -0,0 +1,49 @@ +import typing as t +import construct as cs +from .generic_wrapper import * + + +class TArray( + Adapter[ + t.Any, + t.Any, + ParsedType, + BuildTypes, + ] +): + """ + Adapter for an Array, that transforms the "ListContainer" to an standard "list" while parsing + """ + + # this is unfortunately needed because the stubs are using __new__ instead of __init__ + if t.TYPE_CHECKING: + + def __new__( + cls, + count: ConstantOrContextLambda[int], + subcon: Construct[SubconParsedType, SubconBuildTypes], + discard: bool = False, + ) -> "TArray[t.List[SubconParsedType], t.List[SubconParsedType]]": + ... + + def __init__( + self, + count: ConstantOrContextLambda[int], + subcon: Construct[ParsedType, BuildTypes], + discard: bool = False, + ) -> None: + # init adatper + super(TArray, self).__init__(cs.Array(count, subcon, discard)) # type: ignore + + def _decode( + self, obj: ListContainer[ParsedType], context: Context, path: PathType + ) -> ParsedType: + return list(obj) # type: ignore + + def _encode( + self, + obj: BuildTypes, + context: Context, + path: PathType, + ) -> t.List[BuildTypes]: + return obj # type: ignore diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py new file mode 100644 index 0000000..8e2b718 --- /dev/null +++ b/construct_typed/tenum.py @@ -0,0 +1,68 @@ +import enum +import typing as t +import construct as cs +from .generic_wrapper import * + + +class TEnum(Adapter[int, int, EnumType, t.Union[int, str, EnumType]]): + def __new__( + cls, subcon: Construct[int, int], enum_type: t.Type[EnumType] + ) -> "TEnum[EnumType]": + return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore + + def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): + if not issubclass(enum_type, enum.IntEnum): + raise TypeError( + 'The class "{}" is not an "enum.IntEnum"'.format(enum_type.__name__) + ) + + @classmethod # type: ignore + def _missing_( + cls: t.Type[EnumType], value: t.Union[int, EnumType] + ) -> t.Optional[EnumType]: + if isinstance(value, int): + return cls._create_pseudo_member_(value) # type: ignore + return None # will raise the ValueError in Enum.__new__ + + @classmethod # type: ignore + def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType: + pseudo_member = cls._value2member_map_.get(value, None) # type: ignore + if pseudo_member is None: + new_member = int.__new__(cls, value) # type: ignore + # I expect a name attribute to hold a string, hence str(value) + # However, new_member._name_ = value works, too + new_member._name_ = str(value) + new_member._value_ = value + pseudo_member = cls._value2member_map_.setdefault(value, new_member) # type: ignore + return pseudo_member # type: ignore + + # Monkey-patch the enum type with __missing__ method. So if a enum value + # not found in the enum a new pseudo member is created. + # The idea is taken from: https://stackoverflow.com/a/57179436 + enum_type._missing_ = _missing_ # type: ignore + enum_type._create_pseudo_member_ = _create_pseudo_member_ # type: ignore + + # save enum type + self.enum_type = enum_type + + # init adatper + super(TEnum, self).__init__(subcon) # type: ignore + + def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType: + return self.enum_type(obj) + + def _encode( + self, + obj: t.Union[int, str, EnumType], + context: "cs.Context", + path: "cs.PathType", + ) -> int: + try: + if isinstance(obj, str): + return int(self.enum_type[obj]) + else: + return int(self.enum_type(obj)) + except: + raise cs.MappingError( + "building failed, no mapping for %r" % (obj,), path=path + ) \ No newline at end of file diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py new file mode 100644 index 0000000..bc4c751 --- /dev/null +++ b/construct_typed/tstruct.py @@ -0,0 +1,150 @@ +import typing as t +import textwrap +import dataclasses +from .generic_wrapper import * + + +def StructField( + subcon: Construct[ParsedType, BuildTypes], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, "cs.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) + subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) + + if subcon.flagbuildnone is True: + # some subcons have a predefined default value. all other have "None" + default: t.Any = None + if isinstance(subcon, (cs.Const, cs.Default)): + if callable(subcon.value): # type: ignore + raise ValueError("lamda as default is not supported") + default = subcon.value # type: ignore + + # if subcon builds from "None", set default to "None" + field = dataclasses.field( + default=default, + init=False, + metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, + ) + else: + field = dataclasses.field(metadata={"subcon": subcon}) + + return field # type: ignore + + +class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): + """ + Base class for a typed struct, based on standard dataclasses. + """ + + def __init__( + self, dataclass_type: t.Type[DataclassType], swapped: bool = False + ) -> None: + if not dataclasses.is_dataclass(dataclass_type): + raise TypeError( + 'The class "{}" is not a "dataclasses.dataclass"'.format( + dataclass_type.__name__ + ) + ) + self.dataclass_type = dataclass_type + self.swapped = swapped + + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_type) + if self.swapped: + fields = tuple(reversed(fields)) + + # extract the construct formats from the struct_type + subcon_fields = {} + for field in fields: + subcon_fields[field.name] = field.metadata["subcon"] + + # 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 _decode( + self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType" + ) -> DataclassType: + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_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.dataclass_type(**dc_init) # type: ignore + + # 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) + + return dc + + def _encode( + self, obj: DataclassType, context: "cs.Context", path: "cs.PathType" + ) -> t.Dict[str, t.Any]: + # get all fields from the dataclass + fields = dataclasses.fields(self.dataclass_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 + + +class TStruct(_TStruct[ParsedType, BuildTypes]): + """ + Typed struct, based on standard dataclasses. + """ + + # this is unfortunately needed because the stubs are using __new__ instead of __init__ + if t.TYPE_CHECKING: + + def __new__( + cls, dataclass_type: t.Type[DataclassType], swapped: bool = False + ) -> "TStruct[DataclassType, DataclassType]": + ... + + def _create_subcon( + self, subcon_fields: t.Dict[str, t.Any] + ) -> Construct[t.Any, t.Any]: + return cs.Struct(**subcon_fields) + + +class TBitStruct(_TStruct[ParsedType, BuildTypes]): + """ + Typed bit struct, based on standard dataclasses. + """ + + # this is unfortunately needed because the stubs are using __new__ instead of __init__ + if t.TYPE_CHECKING: + + def __new__( + cls, dataclass_type: t.Type[DataclassType], swapped: bool = False + ) -> "TBitStruct[DataclassType, DataclassType]": + ... + + def _create_subcon( + self, subcon_fields: t.Dict[str, t.Any] + ) -> Construct[t.Any, t.Any]: + return cs.BitStruct(**subcon_fields) diff --git a/construct_typed/tunion.py b/construct_typed/tunion.py new file mode 100644 index 0000000..4c685a4 --- /dev/null +++ b/construct_typed/tunion.py @@ -0,0 +1,44 @@ +import enum +import typing as t +import textwrap +import construct as cs +import dataclasses +from .generic_wrapper import * + + +def UnionField( + subcon: Construct[ParsedType, BuildTypes], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None, +) -> ParsedType: + """ + Create a dataclass field for a "TUnion" 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) + subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed) + + if subcon.flagbuildnone is True: + # some subcons have a predefined default value. all other have "None" + default: t.Any = None + if isinstance(subcon, (cs.Const, cs.Default)): + if callable(subcon.value): # type: ignore + raise ValueError("lamda as default is not supported") + default = subcon.value # type: ignore + + # if subcon builds from "None", set default to "None" + field = dataclasses.field( + default=default, + init=False, + metadata={"subcon": cs.Renamed(subcon, newdocs=doc)}, + ) + else: + field = dataclasses.field(metadata={"subcon": subcon}) + + return field # type: ignore + + +class TUnion(Adapter[t.Any, t.Any, DataclassType, DataclassType]): + pass # TODO \ No newline at end of file diff --git a/setup.py b/setup.py index 93b6242..511f6f5 100644 --- a/setup.py +++ b/setup.py @@ -7,6 +7,7 @@ setup( packages=["construct-stubs", "construct_typed"], package_data={ "construct-stubs": ["*.pyi", "lib/*.pyi"], + "construct_typed": ["py.typed"], }, include_package_data=True, license="MIT", diff --git a/tests/test_typed.py b/tests/test_typed.py index 72a4725..48477dc 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -4,132 +4,151 @@ import enum import dataclasses import typing as t from .declarativeunittest import common, raises -from construct import ( - Int8ub, - Int16ub, - Const, - Pass, - Terminated, - Padding, - SizeofError, - Byte, - Bytes, - Computed, - this, -) -from construct_typed import TStruct, TBitStruct, TSubcon, TEnum +import construct as cs +import construct_typed as cst -def test_typed_struct_1(): +def test_tstruct_1() -> None: @dataclasses.dataclass class TestDataclass: - a: int = TSubcon(Int16ub) - b: int = TSubcon(Int8ub) + a: int = cst.StructField(cs.Int16ub) + b: int = cst.StructField(cs.Int8ub) + + common(cst.TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=1, b=2), 3) + + +def test_tstruct_swapped() -> None: + @dataclasses.dataclass + class TestDataclass: + a: int = cst.StructField(cs.Int16ub) + b: int = cst.StructField(cs.Int8ub) - common(TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=1, b=2), 3) common( - TStruct(TestDataclass, swapped=True), + cst.TStruct(TestDataclass, swapped=True), b"\x02\x00\x01", TestDataclass(a=1, b=2), 3, ) - normal = TStruct(TestDataclass) - swapped = TStruct(TestDataclass, swapped=True) + normal = cst.TStruct(TestDataclass) + swapped = cst.TStruct(TestDataclass, swapped=True) assert str(normal.parse(b"\x00\x01\x02")) == str(swapped.parse(b"\x02\x00\x01")) -def test_typed_struct_2(): + +def test_tstruct_nested() -> None: @dataclasses.dataclass class TestDataclass: @dataclasses.dataclass class InnerDataclass: - b: int = TSubcon(Byte) + b: int = cst.StructField(cs.Byte) - a: InnerDataclass = TSubcon(TStruct(InnerDataclass)) + a: InnerDataclass = cst.StructField(cst.TStruct(InnerDataclass)) common( - TStruct(TestDataclass), + cst.TStruct(TestDataclass), b"\x01", TestDataclass(a=TestDataclass.InnerDataclass(b=1)), 1, ) -def test_typed_struct_3(): + +def test_tstruct_anonymus_fields_1() -> None: @dataclasses.dataclass class TestDataclass: - _1: t.Optional[bytes] = TSubcon(Const(b"\x00")) - _2: None = TSubcon(Padding(1)) - _3: None = TSubcon(Pass) - _4: None = TSubcon(Terminated) + _1: t.Optional[bytes] = cst.StructField(cs.Const(b"\x00")) + _2: None = cst.StructField(cs.Padding(1)) + _3: None = cst.StructField(cs.Pass) + _4: None = cst.StructField(cs.Terminated) common( - TStruct(TestDataclass), + cst.TStruct(TestDataclass), bytes(2), TestDataclass(), - SizeofError, + cs.SizeofError, ) -def test_typed_struct_4(): + +def test_tstruct_anonymus_fields_2() -> None: @dataclasses.dataclass class TestDataclass: - _1: bytes = TSubcon(Bytes(this.missing)) + _1: int = cst.StructField(cs.Computed(7)) + _2: t.Optional[bytes] = cst.StructField(cs.Const(b"JPEG")) + _3: None = cst.StructField(cs.Pass) + _4: None = cst.StructField(cs.Terminated) - assert raises(TStruct(TestDataclass).sizeof) == SizeofError - -def test_typed_struct_5(): - @dataclasses.dataclass - class TestDataclass: - _1: int = TSubcon(Computed(7)) - _2: t.Optional[bytes] = TSubcon(Const(b"JPEG")) - _3: None = TSubcon(Pass) - _4: None = TSubcon(Terminated) - - d = TStruct(TestDataclass) + d = cst.TStruct(TestDataclass) assert d.build(TestDataclass()) == d.build(TestDataclass()) -def test_typed_bit_struct(): +def test_tstruct_missing_dataclass() -> None: + class TestDataclass: + a: int = cst.StructField(cs.Int16ub) + b: int = cst.StructField(cs.Int8ub) + + assert raises(cst.TStruct, TestDataclass) == TypeError + + +def test_tbitstruct() -> None: assert False -def test_enum(): +def test_tenum() -> None: class E(enum.IntEnum): a = 1 b = 2 - a = TEnum(Byte, E) - - common(TEnum(Byte, E), b"\x01", E.a, 1) - common(TEnum(Byte, E), b"\x01", 1, 1) - format = TEnum(Byte, E) + common(cst.TEnum(cs.Byte, E), b"\x01", E.a, 1) + common(cst.TEnum(cs.Byte, E), b"\x01", 1, 1) + format = cst.TEnum(cs.Byte, E) obj = format.parse(b"\x01") assert obj == E.a + assert obj == 1 data = format.build("a") assert data == b"\x01" - common(TEnum(Byte, E), b"\x02", E.b, 1) - common(TEnum(Byte, E), b"\x02", 2, 1) - format = TEnum(Byte, E) + common(cst.TEnum(cs.Byte, E), b"\x02", E.b, 1) + common(cst.TEnum(cs.Byte, E), b"\x02", 2, 1) + format = cst.TEnum(cs.Byte, E) obj = format.parse(b"\x02") assert obj == E.b + assert obj == 2 data = format.build("b") assert data == b"\x02" - common(TEnum(Byte, E), b"\x03", 3, 1) - format = TEnum(Byte, E) + +def test_tenum_missing_value() -> None: + class E(enum.IntEnum): + a = 1 + b = 2 + + common(cst.TEnum(cs.Byte, E), b"\x03", 3, 1) + format = cst.TEnum(cs.Byte, E) obj = format.parse(b"\x03") assert int(obj) == 3 data = format.build(3) assert data == b"\x03" -def test_enum_in_struct(): +def test_tenum_no_int_enum() -> None: + class E(enum.Enum): + a = 1 + b = 2 + + assert raises(cst.TEnum, cs.Byte, E) == TypeError + + +def test_tenum_in_tstruct() -> None: class TestEnum(enum.IntEnum): a = 1 b = 2 @dataclasses.dataclass class TestDataclass: - a: TestEnum = TSubcon(TEnum(Int8ub, TestEnum)) - b: int = TSubcon(Int8ub) + a: TestEnum = cst.StructField(cst.TEnum(cs.Int8ub, TestEnum)) + b: int = cst.StructField(cs.Int8ub) - common(TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=TestEnum.a, b=TestEnum.b), 3) + common( + cst.TStruct(TestDataclass), + b"\x00\x01\x02", + TestDataclass(a=TestEnum.a, b=TestEnum.b), + 3, + )