diff --git a/.vscode/launch.json b/.vscode/launch.json index d94289d..d8f4661 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -9,7 +9,8 @@ "type": "python", "request": "launch", "program": "${file}", - "console": "integratedTerminal" + "console": "integratedTerminal", + "justMyCode": false }, { "name": "Debug Tests", diff --git a/.vscode/settings.json b/.vscode/settings.json index 76bdfe9..2da24ea 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -21,4 +21,9 @@ "reportPrivateUsage": "information", "reportUntypedNamedTuple": "information", }, + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, } \ No newline at end of file diff --git a/README.md b/README.md index b11989c..25f53fb 100644 --- a/README.md +++ b/README.md @@ -68,29 +68,28 @@ A short example: import dataclasses import typing as t from construct import Array, Byte, Const, Int8ub, this -from construct_typed import DataclassMixin, DataclassStruct, EnumBase, TEnum, csfield +from construct_typed import AttrsStruct, Enum, construct, attrs_field -class Orientation(EnumBase): +class Orientation(Enum, constr=Int8ub): # TODO: Implement this HORIZONTAL = 0 VERTICAL = 1 -@dataclasses.dataclass -class Image(DataclassMixin): - signature: bytes = csfield(Const(b"BMP")) - orientation: Orientation = csfield(TEnum(Int8ub, Orientation)) - width: int = csfield(Int8ub) - height: int = csfield(Int8ub) - pixels: t.List[int] = csfield(Array(this.width * this.height, Byte)) +class Image(AttrsStruct): + signature: bytes = attrs_field(Const(b"BMP")) + orientation: Orientation = attrs_field(construct(Orientation)) # TODO: Implement this + width: int = attrs_field(Int8ub) + height: int = attrs_field(Int8ub) + pixels: t.List[int] = attrs_field(Array(this.width * this.height, Byte)) -format = DataclassStruct(Image) +fmt = construct(Image) obj = Image( orientation=Orientation.VERTICAL, width=3, height=2, pixels=[7, 8, 9, 11, 12, 13], ) -print(format.build(obj)) -print(format.parse(b"BMP\x01\x03\x02\x07\x08\t\x0b\x0c\r")) +print(fmt.build(obj)) +print(fmt.parse(b"BMP\x01\x03\x02\x07\x08\t\x0b\x0c\r")) ``` Output: ``` diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index 00d5093..6db86df 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -1,40 +1,34 @@ from .dataclass_struct import ( DataclassBitStruct, - DataclassMixin, DataclassStruct, - TBitStruct, - TContainerBase, - TContainerMixin, - TStruct, - TStructField, csfield, - sfield, ) -from .generic_wrapper import ( +from .attrs_struct import AttrsStruct, attrs_field, this_struct +from .generics import ( Adapter, ConstantOrContextLambda, Construct, Context, ListContainer, PathType, + Constructable, + construct, ) -from .tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum +from .tenum import EnumBase, FlagsEnumBase, EnumConstruct, FlagsEnumConstruct __all__ = [ + "AttrsStruct", + "attrs_field", "DataclassBitStruct", - "DataclassMixin", "DataclassStruct", - "TBitStruct", - "TContainerBase", - "TContainerMixin", - "TStruct", - "TStructField", + "this_struct", "csfield", - "sfield", + "Constructable", + "construct", "EnumBase", "FlagsEnumBase", - "TEnum", - "TFlagsEnum", + "EnumConstruct", + "FlagsEnumConstruct", "Adapter", "ConstantOrContextLambda", "Construct", diff --git a/construct_typed/attrs_struct.py b/construct_typed/attrs_struct.py new file mode 100644 index 0000000..f2778e9 --- /dev/null +++ b/construct_typed/attrs_struct.py @@ -0,0 +1,237 @@ +# -*- coding: utf-8 -*- +# pyright: strict +import textwrap +import typing as t + +import attr +import construct as cs +from .generics import Adapter, Construct, Context, ParsedType, PathType + +T = t.TypeVar("T") + +# Static type inference support via __dataclass_transform__ implemented as per: +# https://github.com/microsoft/pyright/blob/1.1.135/specs/dataclass_transforms.md +def __dataclass_transform__( + *, + eq_default: bool = True, + order_default: bool = False, + kw_only_default: bool = False, + field_descriptors: t.Tuple[t.Union[type, t.Callable[..., t.Any]], ...] = (()), +) -> t.Callable[[T], T]: + return lambda a: a + + +ATTRS_METADATA_KEY = "__construct_typed_subcon" + + +def attrs_field( + subcon: Construct[ParsedType, t.Any], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, +) -> ParsedType: + """ + Helper method for `AttrsStruct` and `AttrsBitStruct` to create the attrs fields. + + This method also processes `Const` and `Default`, to pass these values als default values to the dataclass. + + # TODO: Implement `default` parameter for `attrs_field` + """ + orig_subcon = 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 orig_subcon.flagbuildnone is True: + init = False + default = None + else: + init = True + default = attr.NOTHING + + # Set default values in case of special sucons + if isinstance(orig_subcon, cs.Const): + const_subcon: "cs.Const[t.Any, t.Any, t.Any, t.Any]" = orig_subcon + default = const_subcon.value + elif isinstance(orig_subcon, cs.Default): + default_subcon: "cs.Default[t.Any, t.Any, t.Any, t.Any]" = orig_subcon + if callable(default_subcon.value): + default = None # context lambda is only defined at parsing/building + else: + default = default_subcon.value + + return t.cast( + ParsedType, + attr.field( + default=default, + init=init, + metadata={ATTRS_METADATA_KEY: subcon}, + ), + ) + + +class AttrsConstruct(Adapter[t.Any, t.Any, T, T]): + if t.TYPE_CHECKING: + + def __new__( + cls, + attrs_cls: t.Type[T], + reverse_fields: bool = False, + ) -> "AttrsConstruct[T]": + ... + + def __init__( + self, + attrs_cls: t.Type[T], + reverse_fields: bool = False, + ) -> None: + if not attr.has(attrs_cls): + raise TypeError(f"'{attrs_cls}' has to be a 'attrs' object") + + self.attrs_cls = attrs_cls + self.reverse_fields = reverse_fields + + # get all fields from the dataclass + fields = attr.fields(attrs_cls) + if reverse_fields: + fields = tuple(reversed(fields)) + + # extract the construct formats from the struct_type + subcon_fields = {} + for field in fields: + subcon_fields[field.name] = field.metadata[ATTRS_METADATA_KEY] + + # init adatper + super().__init__(cs.Struct(**subcon_fields)) # type: ignore + + def _decode( + self, + obj: "cs.Container[t.Any]", + context: Context, + path: PathType, + ) -> T: + # get all fields from the dataclass + fields = attr.fields(self.attrs_cls) + + # extract all fields from the container, that are used for create the dataclass object + dc_init = {} + for field in fields: + if field.init: + value = obj[field.name] + dc_init[field.name] = value + + # create object of dataclass + dc = self.attrs_cls(**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 = obj[field.name] + setattr(dc, field.name, value) + + return dc + + def _encode(self, obj: T, context: Context, path: PathType) -> t.Dict[str, t.Any]: + if not isinstance(obj, self.attrs_cls): + raise TypeError(f"'{repr(obj)}' has to be of type {repr(self.attrs_cls)}") + + # get all fields from the dataclass + fields = attr.fields(self.attrs_cls) + + # extract all fields from the container, that are used for create the dataclass object + ret_dict: t.Dict[str, t.Any] = {} + for field in fields: + value = getattr(obj, field.name) + ret_dict[field.name] = value + + return ret_dict + + +# Helper object for defining the `constr` of a `struct`. Will be replaced with the proper construct, when class is created. +this_struct: Construct[t.Any, t.Any] = Construct() + + +def _replace_this_struct(constr: "Construct[t.Any, t.Any]", replacement: t.Any): + """Recursive search for `this_struct` in all SubConstructs and replace it with AttrsStruct""" + subcon = getattr(constr, "subcon", None) + if subcon is this_struct: + setattr(constr, "subcon", replacement) + elif subcon is not None: + _replace_this_struct(subcon, replacement) + else: + raise ValueError( + "Could not find `this_struct`. Only SubConstructs are supported" + ) + + +@__dataclass_transform__(kw_only_default=True, field_descriptors=(attrs_field,)) +class AttrsStruct: + """ + Adapter for a attrs-class for optimised type hints / static autocompletion in comparision to the original Struct. + + Before this construct can be created a dataclasses.dataclass type must be created, which must also derive from DataclassMixin. In this dataclass all fields must be assigned to a construct type using csfield. + + Internally, all fields are converted to a Struct, which does the actual parsing/building. + + Parses to a dataclasses.dataclass instance, and builds from such instance. Size is the sum of all subcon sizes, unless any subcon raises SizeofError. + + Metaclass paramters:: + + :param constr: Create a more complex construct. `this_struct` can be used for representing this AttrsStruct object. + :param reverse_fields: Flag if the fields should be reversed parsed/build + + Example:: + + >>> from construct import Bytes, Int8ub, this + >>> from construct_typed import AttrsStruct, attrs_field, construct + >>> class Image(AttrsStruct): + ... width: int = attrs_field(Int8ub) + ... height: int = attrs_field(Int8ub) + ... pixels: bytes = attrs_field(Bytes(this.height * this.width)) + >>> d = construct(Image) + >>> d.parse(b"\x01\x0212") + Image(width=1, height=2, pixels=b'12') + """ + + @classmethod + def __init_subclass__( + cls, + constr: "cs.Construct[t.Any, t.Any]" = this_struct, + reverse_fields: bool = False, + ): + # validate types + if not isinstance(constr, cs.Construct): # type: ignore + raise ValueError("`constr` parameter has to be an `Construct` object") + if not isinstance(reverse_fields, bool): # type: ignore + raise ValueError("`reverse_fields` parameter has to be an `bool` object") + + # create attrs class + cls = attr.define(cls, kw_only=True, slots=False) + + # create construct format + attrs_constr = AttrsConstruct(cls, reverse_fields) + if constr is this_struct: + constr = attrs_constr + else: + _replace_this_struct(constr, attrs_constr) + + # save construct format and make the class compatible to `Constructable` protocol + setattr(cls, "__construct__", lambda: constr) + + return cls + + # the `construct` library is using the [] access internally, so struct objects + # should also make this possible and not only via the dot access. + def __getitem__(self, key: str) -> t.Any: + return getattr(self, key) + + def __setitem__(self, key: str, value: t.Any) -> None: + setattr(self, key, value) + + if t.TYPE_CHECKING: + + @classmethod + def __construct__(cls: t.Type[T]) -> "AttrsConstruct[T]": + ... diff --git a/construct_typed/dataclass_struct.py b/construct_typed/dataclass_struct.py index 78d2e59..b5ea258 100644 --- a/construct_typed/dataclass_struct.py +++ b/construct_typed/dataclass_struct.py @@ -12,21 +12,249 @@ from construct.lib.containers import ( ) from construct.lib.py3compat import bytestringtype, reprstring, unicodestringtype -from .generic_wrapper import Adapter, Construct, Context, ParsedType, PathType +from construct_typed.generics import Adapter, Construct, Context, ParsedType, PathType + +T = t.TypeVar("T") -class DataclassMixin: +# Static type inference support via __dataclass_transform__ implemented as per: +# https://github.com/microsoft/pyright/blob/1.1.135/specs/dataclass_transforms.md +def __dataclass_transform__( + *, + eq_default: bool = True, + order_default: bool = False, + kw_only_default: bool = False, + field_descriptors: t.Tuple[t.Union[type, t.Callable[..., t.Any]], ...] = (()), +) -> t.Callable[[T], T]: + return lambda a: a + + +DATACLASS_METADATA_KEY = "__construct_typed_subcon" + +if t.TYPE_CHECKING: + # specialisation for constructs, that builds from none and dont have to be declared in the __init__ method + @t.overload + def csfield( + subcon: cs.Construct[ParsedType, None], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, + init: t.Literal[False] = False, + ) -> ParsedType: + ... + + @t.overload + def csfield( + subcon: Construct[ParsedType, t.Any], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, + init: bool = True, + ) -> ParsedType: + ... + + +def csfield( + subcon: Construct[ParsedType, t.Any], + doc: t.Optional[str] = None, + parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, + init: bool = True, +) -> ParsedType: """ - Mixin for the dataclasses which are passed to "DataclassStruct" and "DataclassBitStruct". + Helper method for "DataclassStruct" and "DataclassBitStruct" to create the dataclass fields. - Note: This implementation is different to the 'cs.Container' of the original 'construct' - library. In the original 'cs.Container' some names like "update", "keys", "items", ... can - only accessed via key access (square brackets) and not via attribute access (dot operator), - because they are also method names. This implementation is based on "dataclasses.dataclass" - which only uses modul-level instead of instance-level helper methods.So no instance-level - methods exists and every name can be used. + This method also processes Const and Default, to pass these values als default values to the dataclass. + """ + orig_subcon = 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 orig_subcon.flagbuildnone is True: + init = False + default = None + else: + init = True + default = dataclasses.MISSING + + # Set default values in case of special sucons + if isinstance(orig_subcon, cs.Const): + const_subcon: "cs.Const[t.Any, t.Any, t.Any, t.Any]" = orig_subcon + default = const_subcon.value + elif isinstance(orig_subcon, cs.Default): + default_subcon: "cs.Default[t.Any, t.Any, t.Any, t.Any]" = orig_subcon + if callable(default_subcon.value): + default = None # context lambda is only defined at parsing/building + else: + default = default_subcon.value + + return t.cast( + ParsedType, + dataclasses.field( + default=default, + init=init, + metadata={DATACLASS_METADATA_KEY: subcon}, + ), + ) + + +class DataclassConstruct(Adapter[t.Any, t.Any, T, T]): + """ + Adapter for a dataclasses for optimised type hints / static autocompletion in comparision to the original Struct. + + Before this construct can be created a dataclasses.dataclass type must be created, which must also derive from DataclassMixin. In this dataclass all fields must be assigned to a construct type using csfield. + + Internally, all fields are converted to a Struct, which does the actual parsing/building. + + Parses to a dataclasses.dataclass instance, and builds from such instance. Size is the sum of all subcon sizes, unless any subcon raises SizeofError. + + :param dc_type: Type of the dataclass, which also inherits from DataclassMixin + :param reverse: Flag if the fields of the dataclass should be reversed + + Example:: + + >>> import dataclasses + >>> from construct import Bytes, Int8ub, this + >>> from construct_typed import DataclassMixin, DataclassStruct, csfield, construct + >>> @dataclasses.dataclass + ... class Image(DataclassStruct): + ... width: int = csfield(Int8ub) + ... height: int = csfield(Int8ub) + ... pixels: bytes = csfield(Bytes(this.height * this.width)) + >>> d = construct(Image) + >>> d.parse(b"\x01\x0212") + Image(width=1, height=2, pixels=b'12') """ + subcon: "cs.Struct[t.Any, t.Any]" + if t.TYPE_CHECKING: + + def __new__( + cls, + dc_type: t.Type[T], + reverse: bool = False, + ) -> "DataclassConstruct[T]": + ... + + def __init__( + self, + dc_type: t.Type[T], + reverse: bool = False, + ) -> None: + if not dataclasses.is_dataclass(dc_type): + raise TypeError(f"'{repr(dc_type)}' has to be a 'dataclasses.dataclass'") + self.dc_type = dc_type + self.reverse = reverse + + # get all fields from the dataclass + fields = dataclasses.fields(self.dc_type) + if self.reverse: + fields = tuple(reversed(fields)) + + # extract the construct formats from the struct_type + subcon_fields = {} + for field in fields: + subcon_fields[field.name] = field.metadata[DATACLASS_METADATA_KEY] + + # init adatper + super().__init__(cs.Struct(**subcon_fields)) # type: ignore + + def __getattr__(self, name: str) -> t.Any: + return getattr(self.subcon, name) + + def _decode( + self, obj: "cs.Container[t.Any]", context: Context, path: PathType + ) -> T: + # get all fields from the dataclass + fields = dataclasses.fields(self.dc_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 = obj[field.name] + dc_init[field.name] = value + + # create object of dataclass + dc = self.dc_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 = obj[field.name] + setattr(dc, field.name, value) + + return dc + + def _encode(self, obj: T, context: Context, path: PathType) -> t.Dict[str, t.Any]: + if not isinstance(obj, self.dc_type): + raise TypeError(f"'{repr(obj)}' has to be of type {repr(self.dc_type)}") + + # get all fields from the dataclass + fields = dataclasses.fields(self.dc_type) + + # extract all fields from the container, that are used for create the dataclass object + ret_dict: t.Dict[str, t.Any] = {} + for field in fields: + value = getattr(obj, field.name) + ret_dict[field.name] = value + + return ret_dict + + +# Helper object for defining the `constr` of a `struct`. Will be replaced with the proper construct, when class is created. +this_struct: Construct[t.Any, t.Any] = Construct() + + +def _replace_this_struct(constr: "Construct[t.Any, t.Any]", replacement: t.Any): + """Recursive search for `this_struct` in all SubConstructs and replace it with AttrsStruct""" + subcon = getattr(constr, "subcon", None) + if subcon is this_struct: + setattr(constr, "subcon", replacement) + elif subcon is not None: + _replace_this_struct(subcon, replacement) + else: + raise ValueError( + "Could not find `this_struct`. Only SubConstructs are supported" + ) + + +@__dataclass_transform__(field_descriptors=(csfield,)) +class DataclassStruct: + r""" + TODO: Add Documentation + """ + + @classmethod + def __init_subclass__( + cls, + constr: "cs.Construct[t.Any, t.Any]" = this_struct, + reverse_fields: bool = False, + ): + # validate types + if not isinstance(constr, cs.Construct): # type: ignore + raise ValueError("`constr` parameter has to be an `Construct` object") + if not isinstance(reverse_fields, bool): # type: ignore + raise ValueError("`reverse_fields` parameter has to be an `bool` object") + + # create attrs class + cls = dataclasses.dataclass(cls) + + # create construct format + dc_constr = DataclassConstruct(cls, reverse_fields) + if constr is this_struct: + constr = dc_constr + else: + _replace_this_struct(constr, dc_constr) + + # save construct format and make the class compatible to `Constructable` protocol + setattr(cls, "__construct__", lambda: constr) + + return cls + + # the `construct` library is using the [] access internally, so struct objects + # should also make this possible and not only via the dot access. def __getitem__(self, key: str) -> t.Any: return getattr(self, key) @@ -72,201 +300,42 @@ class DataclassMixin: text.append(indentation.join(str(v).split("\n"))) return "".join(text) - -def csfield( - subcon: Construct[ParsedType, t.Any], - doc: t.Optional[str] = None, - parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, -) -> ParsedType: - """ - Helper method for "DataclassStruct" and "DataclassBitStruct" to create the dataclass fields. - - This method also processes Const and Default, to pass these values als default values to the dataclass. - """ - orig_subcon = 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 orig_subcon.flagbuildnone is True: - init = False - default = None - else: - init = True - default = dataclasses.MISSING - - # Set default values in case of special sucons - if isinstance(orig_subcon, cs.Const): - const_subcon: "cs.Const[t.Any, t.Any, t.Any, t.Any]" = orig_subcon - default = const_subcon.value - elif isinstance(orig_subcon, cs.Default): - default_subcon: "cs.Default[t.Any, t.Any, t.Any, t.Any]" = orig_subcon - if callable(default_subcon.value): - default = None # context lambda is only defined at parsing/building - else: - default = default_subcon.value - - return t.cast( - ParsedType, - dataclasses.field( - default=default, - init=init, - metadata={"subcon": subcon}, - ), - ) - - -DataclassType = t.TypeVar("DataclassType", bound=DataclassMixin) - - -class DataclassStruct(Adapter[t.Any, t.Any, DataclassType, DataclassType]): - """ - Adapter for a dataclasses for optimised type hints / static autocompletion in comparision to the original Struct. - - Before this construct can be created a dataclasses.dataclass type must be created, which must also derive from DataclassMixin. In this dataclass all fields must be assigned to a construct type using csfield. - - Internally, all fields are converted to a Struct, which does the actual parsing/building. - - Parses to a dataclasses.dataclass instance, and builds from such instance. Size is the sum of all subcon sizes, unless any subcon raises SizeofError. - - :param dc_type: Type of the dataclass, which also inherits from DataclassMixin - :param reverse: Flag if the fields of the dataclass should be reversed - - Example:: - - >>> import dataclasses - >>> from construct import Bytes, Int8ub, this - >>> from construct_typed import DataclassMixin, DataclassStruct, csfield - >>> @dataclasses.dataclass - ... class Image(DataclassMixin): - ... width: int = csfield(Int8ub) - ... height: int = csfield(Int8ub) - ... pixels: bytes = csfield(Bytes(this.height * this.width)) - >>> d = DataclassStruct(Image) - >>> d.parse(b"\x01\x0212") - Image(width=1, height=2, pixels=b'12') - """ - - subcon: "cs.Struct[t.Any, t.Any]" if t.TYPE_CHECKING: - def __new__( - cls, - dc_type: t.Type[DataclassType], - reverse: bool = False, - ) -> "DataclassStruct[DataclassType]": + @classmethod + def __construct__(cls: t.Type[T]) -> "DataclassConstruct[T]": ... - def __init__( - self, - dc_type: t.Type[DataclassType], - reverse: bool = False, - ) -> None: - if not issubclass(dc_type, DataclassMixin): - raise TypeError(f"'{repr(dc_type)}' has to be a '{repr(DataclassMixin)}'") - if not dataclasses.is_dataclass(dc_type): - raise TypeError(f"'{repr(dc_type)}' has to be a 'dataclasses.dataclass'") - self.dc_type = dc_type - self.reverse = reverse - # get all fields from the dataclass - fields = dataclasses.fields(self.dc_type) - if self.reverse: - 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().__init__(cs.Struct(**subcon_fields)) # type: ignore - - def __getattr__(self, name: str) -> t.Any: - return getattr(self.subcon, name) - - def _decode( - self, obj: "cs.Container[t.Any]", context: Context, path: PathType - ) -> DataclassType: - # get all fields from the dataclass - fields = dataclasses.fields(self.dc_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 = obj[field.name] - dc_init[field.name] = value - - # create object of dataclass - dc = self.dc_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 = obj[field.name] - setattr(dc, field.name, value) - - return dc - - def _encode( - self, obj: DataclassType, context: Context, path: PathType - ) -> t.Dict[str, t.Any]: - if not isinstance(obj, self.dc_type): - raise TypeError(f"'{repr(obj)}' has to be of type {repr(self.dc_type)}") - - # get all fields from the dataclass - fields = dataclasses.fields(self.dc_type) - - # extract all fields from the container, that are used for create the dataclass object - ret_dict: t.Dict[str, t.Any] = {} - for field in fields: - value = getattr(obj, field.name) - ret_dict[field.name] = value - - return ret_dict - - -def DataclassBitStruct( - dc_type: t.Type[DataclassType], reverse: bool = False -) -> t.Union[ - "cs.Transformed[DataclassType, DataclassType]", - "cs.Restreamed[DataclassType, DataclassType]", -]: +class DataclassBitStruct(DataclassStruct): r""" Makes a DataclassStruct inside a Bitwise. See :class:`~construct.core.Bitwise` and :class:`~construct_typed.dataclass_struct.DatclassStruct` for semantics and raisable exceptions. - :param dc_type: Type of the dataclass, which also inherits from DataclassMixin - :param reverse: Flag if the fields of the dataclass should be reversed + :param constr: TODO + :param reverse_fields: Flag if the fields of the dataclass should be reversed Example:: - DataclassBitStruct <--> Bitwise(DataclassStruct(...)) - >>> import dataclasses + TODO: >>> from construct import BitsInteger, Flag, Nibble, Padding - >>> from construct_typed import DataclassBitStruct, DataclassMixin, csfield - >>> @dataclasses.dataclass - ... class TestDataclass(DataclassMixin): + >>> from construct_typed import DataclassBitStruct, csfield, construct + ... class TestDataclass(DataclassBitStruct): ... a: int = csfield(Flag) ... b: int = csfield(Nibble) ... c: int = csfield(BitsInteger(10)) ... d: None = csfield(Padding(1)) - >>> d = DataclassBitStruct(TestDataclass) + >>> d = construct(TestDataclass) >>> d.parse(b"\x01\x02") TestDataclass(a=False, b=0, c=129, d=None) """ - return cs.Bitwise(DataclassStruct(dc_type, reverse)) - -# support legacy names -TStruct = DataclassStruct -TBitStruct = DataclassBitStruct -TContainerMixin = DataclassMixin -TContainerBase = DataclassMixin -TStructField = csfield -sfield = csfield \ No newline at end of file + @classmethod + def __init_subclass__( + cls, + constr: "cs.Construct[t.Any, t.Any]" = this_struct, + reverse_fields: bool = False, + ): + cls = DataclassStruct.__init_subclass__.__func__(cls, cs.Bitwise(constr), reverse_fields) # type: ignore + return cls diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generics.py similarity index 72% rename from construct_typed/generic_wrapper.py rename to construct_typed/generics.py index aa4af4c..1bb647d 100644 --- a/construct_typed/generic_wrapper.py +++ b/construct_typed/generics.py @@ -39,3 +39,20 @@ else: ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]] PathType = str + + +@t.runtime_checkable +class Constructable(t.Protocol[ParsedType, BuildTypes]): + def __construct__(self) -> "Construct[ParsedType, BuildTypes]": + raise NotImplementedError + + +def construct( + constr: t.Union[ + Constructable[ParsedType, BuildTypes], "Construct[ParsedType, BuildTypes]" + ], +) -> Construct[ParsedType, BuildTypes]: + """Get construct instance of `Constructable` or `Construct`""" + if isinstance(constr, Constructable): + constr = constr.__construct__() + return constr \ No newline at end of file diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index 7c93e33..37eba0c 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -1,17 +1,43 @@ import enum import typing as t -from .generic_wrapper import * +import construct as cs + +from .generics import * + +T = t.TypeVar("T") -# ## TEnum ############################################################################################################ +# ## EnumConstruct ############################################################################################################ class EnumBase(enum.IntEnum): """ - Base class for an Enum used in `construct_typed.TEnum`. + Base class for an Enum used in `construct_typed.EnumConstruct`. This class extends the standard `enum.IntEnum`, so that missing values are automatically generated. """ + @classmethod + def __init_subclass__( + cls, + subcon: "cs.Construct[t.Any, t.Any]", + **kwargs: t.Any, + ): + super().__init_subclass__(**kwargs) + + # validate types + if not isinstance(subcon, cs.Construct): # type: ignore + raise ValueError( + f"`subcon` parameter has to be an `Construct` object but is {type(subcon)}" + ) + + # create construct format + enum_constr = EnumConstruct(subcon, cls) + + # save construct format and make the class compatible to `Constructable` protocol + setattr(cls, "__construct__", lambda: enum_constr) + + return cls + # Extend 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 @@ -33,11 +59,17 @@ class EnumBase(enum.IntEnum): pseudo_member = cls._value2member_map_.setdefault(value, new_member) # type: ignore return pseudo_member # type: ignore + if t.TYPE_CHECKING: + + @classmethod + def __construct__(cls: "t.Type[EnumType]") -> "EnumConstruct[EnumType]": + ... + EnumType = t.TypeVar("EnumType", bound=EnumBase) -class TEnum(Adapter[int, int, EnumType, EnumType]): +class EnumConstruct(Adapter[int, int, EnumType, EnumType]): """ Typed enum. """ @@ -46,7 +78,7 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): def __new__( cls, subcon: Construct[int, int], enum_type: t.Type[EnumType] - ) -> "TEnum[EnumType]": + ) -> "EnumConstruct[EnumType]": ... def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): @@ -59,7 +91,7 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): self.enum_type = t.cast(t.Type[EnumType], enum_type) # type: ignore # init adatper - super(TEnum, self).__init__(subcon) # type: ignore + super(EnumConstruct, self).__init__(subcon) # type: ignore def _decode(self, obj: int, context: Context, path: PathType) -> EnumType: return self.enum_type(obj) @@ -77,15 +109,43 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): ) -# ## TFlagsEnum ####################################################################################################### +# ## FlagsEnumConstruct ####################################################################################################### class FlagsEnumBase(enum.IntFlag): - pass + @classmethod + def __init_subclass__( + cls, + subcon: "cs.Construct[t.Any, t.Any]", + **kwargs: t.Any, + ): + super().__init_subclass__(**kwargs) + + # validate types + if not isinstance(subcon, cs.Construct): # type: ignore + raise ValueError( + f"`subcon` parameter has to be an `Construct` object but is {type(subcon)}" + ) + + # create construct format + enum_constr = FlagsEnumConstruct(subcon, cls) + + # save construct format and make the class compatible to `Constructable` protocol + setattr(cls, "__construct__", lambda: enum_constr) + + return cls + + if t.TYPE_CHECKING: + + @classmethod + def __construct__( + cls: "t.Type[FlagsEnumType]", + ) -> "FlagsEnumConstruct[FlagsEnumType]": + ... FlagsEnumType = t.TypeVar("FlagsEnumType", bound=FlagsEnumBase) -class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): +class FlagsEnumConstruct(Adapter[int, int, FlagsEnumType, FlagsEnumType]): """ Typed enum. """ @@ -94,7 +154,7 @@ class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): def __new__( cls, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType] - ) -> "TFlagsEnum[FlagsEnumType]": + ) -> "FlagsEnumConstruct[FlagsEnumType]": ... def __init__(self, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType]): @@ -107,7 +167,7 @@ class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): self.enum_type = t.cast(t.Type[FlagsEnumType], enum_type) # type: ignore # init adatper - super(TFlagsEnum, self).__init__(subcon) # type: ignore + super(FlagsEnumConstruct, self).__init__(subcon) # type: ignore def _decode(self, obj: int, context: Context, path: PathType) -> FlagsEnumType: return self.enum_type(obj) diff --git a/tests/test_typed.py b/tests/test_typed.py index d74e025..b65a51c 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -1,19 +1,22 @@ # -*- coding: utf-8 -*- # pyright: strict -import dataclasses import enum import typing as t - +import pytest import construct as cs import construct_typed as cst -from construct_typed import DataclassBitStruct, DataclassMixin, DataclassStruct, csfield +from construct_typed import ( + DataclassBitStruct, + DataclassStruct, + csfield, + construct, +) from .declarativeunittest import common, raises, setattrs def test_dataclass_const_default() -> None: - @dataclasses.dataclass - class ConstDefaultTest(DataclassMixin): + class ConstDefaultTest(DataclassStruct): const_bytes: bytes = csfield(cs.Const(b"BMP")) const_int: int = csfield(cs.Const(5, cs.Int8ub)) default_int: int = csfield(cs.Default(cs.Int8ub, 28)) @@ -29,8 +32,7 @@ def test_dataclass_const_default() -> None: def test_dataclass_access() -> None: - @dataclasses.dataclass - class TestTContainer(DataclassMixin): + class TestTContainer(DataclassStruct): a: t.Optional[int] = csfield(cs.Const(1, cs.Byte)) b: int = csfield(cs.Int8ub) @@ -50,17 +52,16 @@ def test_dataclass_access() -> None: assert tcontainer["a"] == 6 # wrong creation - assert raises(lambda: TestTContainer(a=0, b=1)) == TypeError + assert raises(lambda: TestTContainer(a=0, b=1)) == TypeError # type: ignore def test_dataclass_str_repr() -> None: - @dataclasses.dataclass - class Image(DataclassMixin): + class Image(DataclassStruct): signature: t.Optional[bytes] = csfield(cs.Const(b"BMP")) width: int = csfield(cs.Int8ub) height: int = csfield(cs.Int8ub) - format = DataclassStruct(Image) + format = construct(Image) obj = Image(width=3, height=2) assert ( str(obj) @@ -74,64 +75,191 @@ def test_dataclass_str_repr() -> None: def test_dataclass_struct() -> None: - @dataclasses.dataclass - class Image(DataclassMixin): + class Image(DataclassStruct): width: int = csfield(cs.Int8ub) height: int = csfield(cs.Int8ub) pixels: bytes = csfield(cs.Bytes(cs.this.height * cs.this.width)) common( - cst.DataclassStruct(Image), + construct(Image), b"\x01\x0212", Image(width=1, height=2, pixels=b"12"), ) # check __getattr__ - c = cst.DataclassStruct(Image) + c = Image.__construct__() # TODO: construct(Image) assert c.width.name == "width" assert c.height.name == "height" assert c.width.subcon is cs.Int8ub assert c.height.subcon is cs.Int8ub +def test_attrs() -> None: + import attr + + @attr.s(kw_only=True) + class TestAttrs: + a: int = attr.ib() + b: int = attr.ib(default=5) + c: int = attr.ib() + + testattrs1 = TestAttrs(a=5, c=10) + print(testattrs1) + + +import construct_typed.attrs_struct as cst5 + + +def test_attrs_struct_example() -> None: + from construct import Bytes, Int8ub, this + from construct_typed import AttrsStruct, attrs_field, construct + + class Image(AttrsStruct): + width: int = attrs_field(Int8ub) + height: int = attrs_field(Int8ub) + pixels: bytes = attrs_field(Bytes(this.height * this.width)) + + d = construct(Image) + obj = d.parse(b"\x01\x0212") + assert obj.width is obj["width"] + assert obj.height is obj["height"] + assert obj.pixels is obj["pixels"] + + +def test_attrs_struct() -> None: + class Test(cst5.AttrsStruct): + a: int = cst5.attrs_field(cs.Byte) + b: int = cst5.attrs_field(cs.Byte) + c: int = cst5.attrs_field(cs.Byte) + d: int = cst5.attrs_field(cs.Byte) + + common( + cst.construct(Test), + b"\x00\x01\x02\x03", + Test(a=0, b=1, c=2, d=3), + 4, + ) + + +def test_attrs_struct_to_str() -> None: + class Test(cst5.AttrsStruct): + a: int = cst5.attrs_field(cs.Byte) + b: int = cst5.attrs_field(cs.Byte) + c: int = cst5.attrs_field(cs.Byte) + d: int = cst5.attrs_field(cs.Byte) + + obj = Test(a=0, b=1, c=2, d=3) + assert str(obj) == "Test(a=0, b=1, c=2, d=3)" + + +def test_attrs_struct_simple_constr() -> None: + class Test(cst5.AttrsStruct, constr=cst5.this_struct): + a: int = cst5.attrs_field(cs.Byte) + b: int = cst5.attrs_field(cs.Byte) + c: int = cst5.attrs_field(cs.Byte) + d: int = cst5.attrs_field(cs.Byte) + + common( + cst.construct(Test), + b"\x00\x01\x02\x03", + Test(a=0, b=1, c=2, d=3), + 4, + ) + + +def test_attrs_struct_complex_constr() -> None: + class Test(cst5.AttrsStruct, constr=cs.Bitwise(cst5.this_struct)): + a: int = cst5.attrs_field(cs.BitsInteger(2)) + b: int = cst5.attrs_field(cs.BitsInteger(2)) + c: int = cst5.attrs_field(cs.BitsInteger(2)) + d: int = cst5.attrs_field(cs.BitsInteger(2)) + + common( + cst.construct(Test), + b"\x1b", + Test(a=0, b=1, c=2, d=3), + 1, + ) + + +def test_attrs_struct_overloaded_attributes() -> None: + class Test(cst5.AttrsStruct): + a: int = cst5.attrs_field(cs.Byte) + b: int = cst5.attrs_field(cs.Byte) + subcon: int = cst5.attrs_field( + cs.Byte + ) # this is also an attribute from Construct + docs: int = cst5.attrs_field( + cs.Byte + ) # this is also an attribute from Construct + + common( + cst.construct(Test), + b"\x00\x01\x02\x03", + Test(a=0, b=1, subcon=2, docs=3), + 4, + ) + + +def test_attrs_struct_reverse_fields() -> None: + class Test(cst5.AttrsStruct, reverse_fields=True): + a: int = cst5.attrs_field(cs.Byte) + b: int = cst5.attrs_field(cs.Byte) + c: int = cst5.attrs_field(cs.Byte) + d: int = cst5.attrs_field(cs.Byte) + + common( + cst.construct(Test), + b"\x03\x02\x01\x00", + Test(a=0, b=1, c=2, d=3), + 4, + ) + + +def test_attrs_struct_unsupported_param() -> None: + with pytest.raises(TypeError): + + class Test(cst5.AttrsStruct, strange_parameter=True): # type: ignore + a: int = cst5.attrs_field(cs.Byte) + + +@pytest.mark.skip +def test_attrs_default() -> None: + # TODO: Implement `default` parameter for `attrs_field` + raise NotImplementedError + + def test_dataclass_struct_reverse() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainerReverse(DataclassStruct, reverse_fields=True): a: int = csfield(cs.Int16ub) b: int = csfield(cs.Int8ub) common( - DataclassStruct(TestContainer, reverse=True), + cst.construct(TestContainerReverse), b"\x02\x00\x01", - TestContainer(a=1, b=2), + TestContainerReverse(a=1, b=2), 3, ) - normal = DataclassStruct(TestContainer) - reverse = DataclassStruct(TestContainer, reverse=True) - assert str(normal.parse(b"\x00\x01\x02")) == str(reverse.parse(b"\x02\x00\x01")) def test_dataclass_struct_nested() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): - @dataclasses.dataclass - class InnerDataclass(DataclassMixin): + class TestContainer(DataclassStruct): + class InnerDataclass(DataclassStruct): b: int = csfield(cs.Byte) c: bytes = csfield(cs.Bytes(cs.this._.length)) length: int = csfield(cs.Byte) - a: InnerDataclass = csfield(DataclassStruct(InnerDataclass)) + a: InnerDataclass = csfield(cst.construct(InnerDataclass)) common( - DataclassStruct(TestContainer), + cst.construct(TestContainer), b"\x02\x01\xF1\xF2", TestContainer(length=2, a=TestContainer.InnerDataclass(b=1, c=b"\xF1\xF2")), ) def test_dataclass_struct_default_field() -> None: - @dataclasses.dataclass - class Image(DataclassMixin): + class Image(DataclassStruct): width: int = csfield(cs.Int8ub) height: int = csfield(cs.Int8ub) pixels: t.Optional[bytes] = csfield( @@ -142,7 +270,7 @@ def test_dataclass_struct_default_field() -> None: ) common( - DataclassStruct(Image), + cst.construct(Image), b"\x02\x03\x00\x00\x00\x00\x00\x00", setattrs(Image(2, 3), pixels=bytes(6)), sample_building=Image(2, 3), @@ -150,12 +278,11 @@ def test_dataclass_struct_default_field() -> None: def test_dataclass_struct_const_field() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): const_field: t.Optional[bytes] = csfield(cs.Const(b"\x00")) common( - DataclassStruct(TestContainer), + cst.construct(TestContainer), bytes(1), setattrs(TestContainer(), const_field=b"\x00"), 1, @@ -163,7 +290,7 @@ def test_dataclass_struct_const_field() -> None: assert ( raises( - DataclassStruct(TestContainer).build, + cst.construct(TestContainer).build, setattrs(TestContainer(), const_field=b"\x01"), ) == cs.ConstError @@ -171,12 +298,11 @@ def test_dataclass_struct_const_field() -> None: def test_dataclass_struct_array_field() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): array_field: t.List[int] = csfield(cs.Array(5, cs.Int8ub)) common( - DataclassStruct(TestContainer), + cst.construct(TestContainer), bytes(5), TestContainer(array_field=[0, 0, 0, 0, 0]), 5, @@ -184,15 +310,14 @@ def test_dataclass_struct_array_field() -> None: def test_dataclass_struct_anonymus_fields_1() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): _1: t.Optional[bytes] = csfield(cs.Const(b"\x00")) _2: None = csfield(cs.Padding(1)) _3: None = csfield(cs.Pass) _4: None = csfield(cs.Terminated) common( - DataclassStruct(TestContainer), + cst.construct(TestContainer), bytes(2), setattrs(TestContainer(), _1=b"\x00"), cs.SizeofError, @@ -200,22 +325,20 @@ def test_dataclass_struct_anonymus_fields_1() -> None: def test_dataclass_struct_anonymus_fields_2() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): _1: int = csfield(cs.Computed(7)) _2: t.Optional[bytes] = csfield(cs.Const(b"JPEG")) _3: None = csfield(cs.Pass) _4: None = csfield(cs.Terminated) - d = DataclassStruct(TestContainer) + d = cst.construct(TestContainer) assert d.build(TestContainer()) == d.build(TestContainer()) def test_dataclass_struct_overloaded_method() -> None: # Test dot access to some names that are not accessable via dot # in the original 'cs.Container'. - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): clear: int = csfield(cs.Int8ul) copy: int = csfield(cs.Int8ul) fromkeys: int = csfield(cs.Int8ul) @@ -231,7 +354,7 @@ def test_dataclass_struct_overloaded_method() -> None: update: int = csfield(cs.Int8ul) values: int = csfield(cs.Int8ul) - d = DataclassStruct(TestContainer) + d = construct(TestContainer) obj = d.parse( d.build( TestContainer( @@ -268,44 +391,22 @@ def test_dataclass_struct_overloaded_method() -> None: assert obj.values == 14 -def test_dataclass_struct_no_dataclass() -> None: - class TestContainer(DataclassMixin): - a: int = csfield(cs.Int16ub) - b: int = csfield(cs.Int8ub) - - assert raises(lambda: DataclassStruct(TestContainer)) == TypeError - - -def test_dataclass_struct_no_DataclassMixin() -> None: - @dataclasses.dataclass - class TestContainer: - a: int = csfield(cs.Int16ub) - b: int = csfield(cs.Int8ub) - - cls = t.cast(t.Type[DataclassMixin], TestContainer) - assert raises(lambda: DataclassStruct(cls)) == TypeError - - def test_dataclass_struct_wrong_container() -> None: - @dataclasses.dataclass - class TestContainer1(DataclassMixin): + class TestContainer1(DataclassStruct): a: int = csfield(cs.Int16ub) b: int = csfield(cs.Int8ub) - @dataclasses.dataclass - class TestContainer2(DataclassMixin): + class TestContainer2(DataclassStruct): a: int = csfield(cs.Int16ub) b: int = csfield(cs.Int8ub) assert ( - raises(DataclassStruct(TestContainer1).build, TestContainer2(a=1, b=2)) - == TypeError + raises(construct(TestContainer1).build, TestContainer2(a=1, b=2)) == TypeError ) def test_dataclass_struct_doc() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassStruct): a: int = csfield(cs.Int16ub, "This is the documentation of a") b: int = csfield( cs.Int8ub, doc="This is the documentation of b\nwhich is multiline" @@ -318,7 +419,7 @@ def test_dataclass_struct_doc() -> None: """, ) - format = DataclassStruct(TestContainer) + format = TestContainer.__construct__() # TODO: construct(TestContainer) common(format, b"\x00\x01\x02\x03", TestContainer(a=1, b=2, c=3), 4) assert format.subcon.a.docs == "This is the documentation of a" @@ -330,39 +431,36 @@ def test_dataclass_struct_doc() -> None: def test_dataclass_bitstruct() -> None: - @dataclasses.dataclass - class TestContainer(DataclassMixin): + class TestContainer(DataclassBitStruct): a: int = csfield(cs.BitsInteger(7)) b: int = csfield(cs.Bit) c: int = csfield(cs.BitsInteger(8)) - print("") - common( - DataclassBitStruct(TestContainer), + construct(TestContainer), b"\xFD\x12", TestContainer(a=0x7E, b=1, c=0x12), 2, ) # check __getattr__ - c = DataclassStruct(TestContainer) - assert c.a.name == "a" - assert c.b.name == "b" - assert c.c.name == "c" - assert isinstance(c.a.subcon, cs.BitsInteger) - assert c.b.subcon is cs.Bit - assert isinstance(c.c.subcon, cs.BitsInteger) + c = TestContainer.__construct__() + assert c.subcon.a.name == "a" + assert c.subcon.b.name == "b" + assert c.subcon.c.name == "c" + assert isinstance(c.subcon.a.subcon, cs.BitsInteger) + assert c.subcon.b.subcon is cs.Bit + assert isinstance(c.subcon.c.subcon, cs.BitsInteger) def test_tenum() -> None: - class TestEnum(cst.EnumBase): + class TestEnum(cst.EnumBase, subcon=cs.Byte): one = 1 two = 2 four = 4 eight = 8 - d = cst.TEnum(cs.Byte, TestEnum) + d = cst.construct(TestEnum) common(d, b"\x01", TestEnum.one, 1) common(d, b"\xff", TestEnum(255), 1) @@ -381,51 +479,46 @@ def test_tenum_no_enumbase() -> None: b = 2 cls = t.cast(t.Type[cst.EnumBase], E) - assert raises(lambda: cst.TEnum(cs.Byte, cls)) == TypeError + assert raises(lambda: cst.EnumConstruct(cs.Byte, cls)) == TypeError -def test_dataclass_struct_wrong_enumbase() -> None: - class E1(cst.EnumBase): - a = 1 - b = 2 +def test_tenum_no_subcon() -> None: + with pytest.raises(TypeError): - class E2(cst.EnumBase): - a = 1 - b = 2 - - assert raises(cst.TEnum(cs.Byte, E1).build, E2.a) == TypeError + class E1(cst.EnumBase): # type: ignore + a = 1 + b = 2 def test_tenum_in_tstruct() -> None: - class TestEnum(cst.EnumBase): + class TestEnum(cst.EnumBase, subcon=cs.Int8ub): a = 1 b = 2 - @dataclasses.dataclass - class TestContainer(DataclassMixin): - a: TestEnum = csfield(cst.TEnum(cs.Int8ub, TestEnum)) + class TestContainer(DataclassStruct): + a: TestEnum = csfield(cst.construct(TestEnum)) b: int = csfield(cs.Int8ub) common( - DataclassStruct(TestContainer), + construct(TestContainer), b"\x01\x02", TestContainer(a=TestEnum.a, b=2), 2, ) assert ( - raises(cst.TEnum(cs.Byte, TestEnum).build, TestContainer(a=1, b=2)) == TypeError # type: ignore + raises(cst.construct(TestEnum).build, TestContainer(a=1, b=2)) == TypeError # type: ignore ) def test_tenum_flags() -> None: - class TestEnum(cst.FlagsEnumBase): + class TestEnum(cst.FlagsEnumBase, subcon=cs.Byte): one = 1 two = 2 four = 4 eight = 8 - d = cst.TFlagsEnum(cs.Byte, TestEnum) + d = cst.construct(TestEnum) common(d, b"\x03", TestEnum.one | TestEnum.two, 1) assert d.build(TestEnum(0)) == b"\x00" assert d.build(TestEnum.one | TestEnum.two) == b"\x03"