wrapper/construct-typing
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

added "TContainerBase"

Browse source
This commit is contained in:
Tim Rid 2021-01-02 14:49:34 +01:00
parent 7d22a9aff4
commit ffe3211326
4 changed files with 182 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

7
construct_typed/__init__.py
View file

@ -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",
]

4
construct_typed/generic_wrapper.py
View file

@ -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
PathType = str

79
construct_typed/tstruct.py
View file

@ -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.
"""

147
tests/test_typed.py
View file

@ -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
)