From 9c7df9f8e39bc37c42ac7007d7fa42e0fde721da Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Sun, 9 Jan 2022 21:22:36 +0100 Subject: [PATCH] first working implementation of `AttrsStruct` --- README.md | 23 +- construct_typed/__init__.py | 12 +- construct_typed/attrs_struct.py | 248 ++++++++++++++++++ construct_typed/dataclass_struct.py | 2 +- .../{generic_wrapper.py => generics.py} | 17 ++ construct_typed/tenum.py | 4 +- tests/test_typed.py | 131 ++++++++- 7 files changed, 421 insertions(+), 16 deletions(-) create mode 100644 construct_typed/attrs_struct.py rename construct_typed/{generic_wrapper.py => generics.py} (72%) 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..092650d 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -10,17 +10,25 @@ from .dataclass_struct import ( csfield, sfield, ) -from .generic_wrapper import ( +from .attrs_struct import ( + AttrsStruct, + attrs_field +) +from .generics import ( Adapter, ConstantOrContextLambda, Construct, Context, ListContainer, PathType, + Constructable, + construct ) from .tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum __all__ = [ + "AttrsStruct", + "attrs_field", "DataclassBitStruct", "DataclassMixin", "DataclassStruct", @@ -30,6 +38,8 @@ __all__ = [ "TStruct", "TStructField", "csfield", + "Constructable", + "construct", "sfield", "EnumBase", "FlagsEnumBase", diff --git a/construct_typed/attrs_struct.py b/construct_typed/attrs_struct.py new file mode 100644 index 0000000..9e2ff7e --- /dev/null +++ b/construct_typed/attrs_struct.py @@ -0,0 +1,248 @@ +# -*- coding: utf-8 -*- +# pyright: strict +import textwrap +import typing as t + +import attr +import construct as cs +import abc +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 AttrsStructMeta(abc.ABCMeta): + def __new__( + metacls, # type: ignore + name: str, + bases: t.Tuple[type, ...], + namespace: t.Dict[str, t.Any], + **kwargs: t.Any, + ): + # extract parameters from kwargs + constr: "cs.Construct[t.Any, t.Any]" = kwargs.pop("constr", this_struct) + if not isinstance(constr, cs.Construct): # type: ignore + raise ValueError("`constr` parameter has to be an `Construct` object") + reverse_fields = kwargs.pop("reverse_fields", False) + if not isinstance(reverse_fields, bool): + raise ValueError("`reverse_fields` parameter has to be an `bool` object") + if len(kwargs) > 0: + unsupp_parm = ", ".join([f"'{k}'" for k in kwargs.keys()]) + raise ValueError(f"unsupported parameter(s) detected: {unsupp_parm}") + + # create new class object + cls = super().__new__(metacls, name, bases, namespace) + + # create attrs class + cls = attr.define(cls, kw_only=True, slots=False) + + # create construct format + attrs_constr = AttrsConstruct(cls, reverse_fields) # type: ignore + 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 + + if t.TYPE_CHECKING: + + def __construct__(self: t.Type[T]) -> "AttrsConstruct[T]": + ... + + +class AttrsStruct(metaclass=AttrsStructMeta): + """ + 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') + """ + + pass diff --git a/construct_typed/dataclass_struct.py b/construct_typed/dataclass_struct.py index 78d2e59..f22559a 100644 --- a/construct_typed/dataclass_struct.py +++ b/construct_typed/dataclass_struct.py @@ -12,7 +12,7 @@ from construct.lib.containers import ( ) from construct.lib.py3compat import bytestringtype, reprstring, unicodestringtype -from .generic_wrapper import Adapter, Construct, Context, ParsedType, PathType +from .generics import Adapter, Construct, Context, ParsedType, PathType class DataclassMixin: 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..ebae090 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -1,7 +1,7 @@ import enum import typing as t -from .generic_wrapper import * +from .generics import * # ## TEnum ############################################################################################################ @@ -33,6 +33,8 @@ class EnumBase(enum.IntEnum): pseudo_member = cls._value2member_map_.setdefault(value, new_member) # type: ignore return pseudo_member # type: ignore + # TODO: Add `__construct__` method to support `Constructable` protocol + EnumType = t.TypeVar("EnumType", bound=EnumBase) diff --git a/tests/test_typed.py b/tests/test_typed.py index d74e025..72506eb 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -3,7 +3,7 @@ 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 @@ -94,6 +94,135 @@ def test_dataclass_struct() -> None: 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) + d.parse(b"\x01\x0212") + + +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(ValueError, match=r".*strange_parameter.*"): + + class Test(cst5.AttrsStruct, strange_parameter=True): # type: ignore + a: int = cst5.attrs_field(cs.Byte) + +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):