From ffe3211326ea4672358abfba1819682199857b8a Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Sat, 2 Jan 2021 14:49:34 +0100 Subject: [PATCH] added "TContainerBase" --- construct_typed/__init__.py | 7 +- construct_typed/generic_wrapper.py | 4 +- construct_typed/tstruct.py | 79 ++++++++++++---- tests/test_typed.py | 147 ++++++++++++++++++++++------- 4 files changed, 182 insertions(+), 55 deletions(-) diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index 4de653b..0491326 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -7,8 +7,8 @@ from .generic_wrapper import ( PathType, ) from .tarray import TArray -from .tenum import TEnum, EnumBase, TFlagsEnum, FlagsEnumBase -from .tstruct import TBitStruct, TStruct, TStructField +from .tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum +from .tstruct import TBitStruct, TStruct, TStructField, TContainerBase from .tunion import TUnion, TUnionField __all__ = [ @@ -23,9 +23,10 @@ __all__ = [ "Construct", "Adapter", "ListContainer", + "TContainerBase", "Context", "ConstantOrContextLambda", "PathType", "TFlagsEnum", - "FlagsEnumBase" + "FlagsEnumBase", ] diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generic_wrapper.py index 2f4685f..c76974c 100644 --- a/construct_typed/generic_wrapper.py +++ b/construct_typed/generic_wrapper.py @@ -21,6 +21,8 @@ if t.TYPE_CHECKING: from construct import Context as Context from construct import ListContainer as ListContainer 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): @@ -39,4 +41,4 @@ else: pass ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]] - PathType = str \ No newline at end of file + PathType = str diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py index 6196ed5..f62e99c 100644 --- a/construct_typed/tstruct.py +++ b/construct_typed/tstruct.py @@ -4,13 +4,45 @@ import typing as t from .generic_wrapper import * -DataclassType = t.TypeVar("DataclassType") + +if t.TYPE_CHECKING: + + class TContainerBase(cs.Container[t.Any]): + def __init__(self, *args: t.Any, **kwargs: t.Any): + ... + + +else: + + class TContainerBase(cs.Container): + def __init__(self, *args, **kwargs): + raise RuntimeError( + "this should never be called, because it shoult be overwritten by 'dataclasses.dataclass'" + ) + + # if accessing via an field via dot access, return the object from the dict + def __getattribute__(self, name): + if name in self: + return self[name] + else: + return super().__getattribute__(name) + + def __post_init__(self): + # 1. fix the __keys_order__ of the cs.Container + # 2. append fields with init=False to the dict of the cs.Container + self.__keys_order__ = [] + for field in dataclasses.fields(self): + value = getattr(self, field.name) + if field.init is True: + self.__keys_order__.append(field.name) + else: + self[field.name] = value def TStructField( subcon: Construct[ParsedType, BuildTypes], doc: t.Optional[str] = None, - parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None, + parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None, ) -> ParsedType: """ Create a dataclass field for a "TStruct" and "TBitStruct" from a subcon. @@ -34,23 +66,32 @@ def TStructField( return field # type: ignore -class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): +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, dataclass_type: t.Type[ParsedType], swapped: bool = False + self, container_type: t.Type[ContainerType], swapped: bool = False ) -> None: - if not dataclasses.is_dataclass(dataclass_type): + if not issubclass(container_type, TContainerBase): raise TypeError( - "'{}' has to be a 'dataclasses.dataclass'".format(repr(dataclass_type)) + "'{}' has to be a '{}'".format( + repr(container_type), repr(TContainerBase) + ) ) - self.dataclass_type = dataclass_type + 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 # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_type) + fields = dataclasses.fields(self.container_type) if self.swapped: fields = tuple(reversed(fields)) @@ -68,10 +109,10 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): raise NotImplementedError def _decode( - self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType" - ) -> ParsedType: + self, obj: "cs.Container[t.Any]", context: Context, path: PathType + ) -> ContainerType: # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_type) + fields = dataclasses.fields(self.container_type) # extract all fields from the container, that are used for create the dataclass object dc_init = {} @@ -81,7 +122,7 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): dc_init[field.name] = value # create object of dataclass - dc = self.dataclass_type(**dc_init) # type: ignore + dc = self.container_type(**dc_init) # extract all other values from the container, an pass it to the dataclass for field in fields: @@ -92,11 +133,11 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): return dc def _encode( - self, obj: BuildTypes, context: "cs.Context", path: "cs.PathType" + self, obj: BuildTypes, context: Context, path: PathType ) -> t.Dict[str, t.Any]: - if isinstance(obj, self.dataclass_type): + if isinstance(obj, self.container_type): # get all fields from the dataclass - fields = dataclasses.fields(self.dataclass_type) + fields = dataclasses.fields(self.container_type) # extract all fields from the container, that are used for create the dataclass object ret_dict = {} @@ -105,10 +146,12 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): ret_dict[field.name] = value return ret_dict - raise TypeError("'{}' has to be of type {}".format(repr(obj), repr(self.dataclass_type))) + raise TypeError( + "'{}' has to be of type {}".format(repr(obj), repr(self.container_type)) + ) -class TStruct(_TStruct[ParsedType, ParsedType]): +class TStruct(_TStruct[ContainerType, ContainerType]): """ Typed struct, based on standard dataclasses. """ @@ -119,7 +162,7 @@ class TStruct(_TStruct[ParsedType, ParsedType]): return cs.Struct(**subcon_fields) -class TBitStruct(_TStruct[ParsedType, ParsedType]): +class TBitStruct(_TStruct[ContainerType, ContainerType]): """ Typed bit struct, based on standard dataclasses. """ diff --git a/tests/test_typed.py b/tests/test_typed.py index 2adbd52..921154a 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -10,52 +10,124 @@ import construct_typed as cst from .declarativeunittest import common, raises, setattrs +def test_tcontainer_compare_with_dataclass(): + @dataclasses.dataclass + class TestContainer: + a: t.Optional[int] = cst.TStructField(cs.Const(1, cs.Byte)) + b: int = cst.TStructField(cs.Int8ub) + + @dataclasses.dataclass + class TestTContainer(cst.TContainerBase): + a: t.Optional[int] = cst.TStructField(cs.Const(1, cs.Byte)) + b: int = cst.TStructField(cs.Int8ub) + + datacls = TestContainer(b=1) + tcontainer = TestTContainer(b=1) + + # ##### compare dot & dict access ##### + # dataclass + assert datacls.a == None + assert raises(lambda: datacls["a"]) == TypeError # type: ignore + assert datacls.b == 1 + assert raises(lambda: datacls["b"]) == TypeError # type: ignore + + datacls.a = 5 + assert datacls.a == 5 + assert raises(lambda: datacls["a"]) == TypeError # type: ignore + try: + datacls["a"] = 5 # type: ignore + except Exception as e: + assert e.__class__ == TypeError + + # tcontainer + assert tcontainer.a == None + assert tcontainer["a"] == None + assert tcontainer.b == 1 + assert tcontainer["b"] == 1 + + tcontainer.a = 5 + assert tcontainer.a == 5 + assert tcontainer["a"] == 5 + tcontainer["a"] = 5 + assert tcontainer.a == 5 + assert tcontainer["a"] == 5 + + # ##### compare fields ##### + for datacls_field, tcontainer_field in zip( + dataclasses.fields(TestContainer), dataclasses.fields(TestTContainer) + ): + assert repr(datacls_field) == repr(tcontainer_field) + + # ##### compare wrong creation ##### + assert raises(lambda: TestContainer(a=0, b=1)) == TypeError + assert raises(lambda: TestTContainer(a=0, b=1)) == TypeError + + +def test_tcontainer_order() -> None: + @dataclasses.dataclass + class Image(cst.TContainerBase): + signature: t.Optional[bytes] = cst.TStructField(cs.Const(b"BMP")) + width: int = cst.TStructField(cs.Int8ub) + height: int = cst.TStructField(cs.Int8ub) + + format = cst.TStruct(Image) + obj = Image(width=3, height=2) + assert ( + str(obj) == "Container: \n signature = None\n width = 3\n height = 2" + ) + obj = format.parse(format.build(obj)) + assert ( + str(obj) + == "Container: \n signature = b'BMP' (total 3)\n width = 3\n height = 2" + ) + + def test_tstruct() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) - common(cst.TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=1, b=2), 3) + common(cst.TStruct(TestContainer), b"\x00\x01\x02", TestContainer(a=1, b=2), 3) def test_tstruct_swapped() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) common( - cst.TStruct(TestDataclass, swapped=True), + cst.TStruct(TestContainer, swapped=True), b"\x02\x00\x01", - TestDataclass(a=1, b=2), + TestContainer(a=1, b=2), 3, ) - normal = cst.TStruct(TestDataclass) - swapped = cst.TStruct(TestDataclass, swapped=True) + normal = cst.TStruct(TestContainer) + swapped = cst.TStruct(TestContainer, swapped=True) assert str(normal.parse(b"\x00\x01\x02")) == str(swapped.parse(b"\x02\x00\x01")) def test_tstruct_nested() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): @dataclasses.dataclass - class InnerDataclass: + class InnerDataclass(cst.TContainerBase): b: int = cst.TStructField(cs.Byte) a: InnerDataclass = cst.TStructField(cst.TStruct(InnerDataclass)) common( - cst.TStruct(TestDataclass), + cst.TStruct(TestContainer), b"\x01", - TestDataclass(a=TestDataclass.InnerDataclass(b=1)), + TestContainer(a=TestContainer.InnerDataclass(b=1)), 1, ) def test_tstruct_default_field() -> None: @dataclasses.dataclass - class Image: + class Image(cst.TContainerBase): width: int = cst.TStructField(cs.Int8ub) height: int = cst.TStructField(cs.Int8ub) pixels: t.Optional[bytes] = cst.TStructField( @@ -75,20 +147,20 @@ def test_tstruct_default_field() -> None: def test_tstruct_const_field() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): const_field: t.Optional[bytes] = cst.TStructField(cs.Const(b"\x00")) common( - cst.TStruct(TestDataclass), + cst.TStruct(TestContainer), bytes(1), - setattrs(TestDataclass(), const_field=b"\x00"), + setattrs(TestContainer(), const_field=b"\x00"), 1, ) assert ( raises( - cst.TStruct(TestDataclass).build, - setattrs(TestDataclass(), const_field=b"\x01"), + cst.TStruct(TestContainer).build, + setattrs(TestContainer(), const_field=b"\x01"), ) == cs.ConstError ) @@ -96,53 +168,62 @@ def test_tstruct_const_field() -> None: def test_tstruct_anonymus_fields_1() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): _1: t.Optional[bytes] = cst.TStructField(cs.Const(b"\x00")) _2: None = cst.TStructField(cs.Padding(1)) _3: None = cst.TStructField(cs.Pass) _4: None = cst.TStructField(cs.Terminated) common( - cst.TStruct(TestDataclass), + cst.TStruct(TestContainer), bytes(2), - setattrs(TestDataclass(), _1=b"\x00"), + setattrs(TestContainer(), _1=b"\x00"), cs.SizeofError, ) def test_tstruct_anonymus_fields_2() -> None: @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): _1: int = cst.TStructField(cs.Computed(7)) _2: t.Optional[bytes] = cst.TStructField(cs.Const(b"JPEG")) _3: None = cst.TStructField(cs.Pass) _4: None = cst.TStructField(cs.Terminated) - d = cst.TStruct(TestDataclass) - assert d.build(TestDataclass()) == d.build(TestDataclass()) + d = cst.TStruct(TestContainer) + assert d.build(TestContainer()) == d.build(TestContainer()) def test_tstruct_no_dataclass() -> None: - class TestDataclass: + class TestContainer(cst.TContainerBase): a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) - assert raises(lambda: cst.TStruct(TestDataclass)) == TypeError + assert raises(lambda: cst.TStruct(TestContainer)) == TypeError -def test_tstruct_wrong_dataclass() -> None: +def test_tstruct_no_tcontainerbase() -> None: @dataclasses.dataclass - class TestDataclass1: + class TestContainer: + a: int = cst.TStructField(cs.Int16ub) + b: int = cst.TStructField(cs.Int8ub) + + assert raises(lambda: cst.TStruct(TestContainer)) == TypeError + + +def test_tstruct_wrong_container() -> None: + @dataclasses.dataclass + class TestContainer1(cst.TContainerBase): a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) @dataclasses.dataclass - class TestDataclass2: + class TestContainer2(cst.TContainerBase): a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) assert ( - raises(cst.TStruct(TestDataclass1).build, TestDataclass2(a=1, b=2)) == TypeError + raises(cst.TStruct(TestContainer1).build, TestContainer2(a=1, b=2)) == TypeError ) @@ -196,19 +277,19 @@ def test_tenum_in_tstruct() -> None: b = 2 @dataclasses.dataclass - class TestDataclass: + class TestContainer(cst.TContainerBase): a: TestEnum = cst.TStructField(cst.TEnum(cs.Int8ub, TestEnum)) b: int = cst.TStructField(cs.Int8ub) common( - cst.TStruct(TestDataclass), + cst.TStruct(TestContainer), b"\x01\x02", - TestDataclass(a=TestEnum.a, b=2), + TestContainer(a=TestEnum.a, b=2), 2, ) assert ( - raises(cst.TEnum(cs.Byte, TestEnum).build, TestDataclass(a=1, b=2)) == TypeError # type: ignore + raises(cst.TEnum(cs.Byte, TestEnum).build, TestContainer(a=1, b=2)) == TypeError # type: ignore )