From def9e93da7913ba0f22be6e89fe29a6681adb002 Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Sat, 2 Jan 2021 01:12:04 +0100 Subject: [PATCH] - enhanced TEnum - added EnumBase - extended README.md --- README.md | 82 ++++++++++++++++++++------- construct_typed/__init__.py | 3 +- construct_typed/generic_wrapper.py | 2 +- construct_typed/tenum.py | 91 ++++++++++++++---------------- construct_typed/tstruct.py | 4 +- tests/test_typed.py | 63 ++++++++++++--------- 6 files changed, 144 insertions(+), 101 deletions(-) diff --git a/README.md b/README.md index 413d53f..f03a3bb 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # construct-typing -This project is an extension of the python package *construct*. This Repository consitst of two packages: +This project is an extension of the python package [*construct*](https://pypi.org/project/construct/). This Repository consists of two packages: -- **construct-stubs**: Adding .pyi for the whole *construct* package (according to [PEP 561 stub-only packages](https://www.python.org/dev/peps/pep-0561/#stub-only-packages)) +- **construct-stubs**: Adding .pyi for the whole *construct 2.10* package (according to [PEP 561 stub-only packages](https://www.python.org/dev/peps/pep-0561/#stub-only-packages)) - **construct_typed**: Adding the additional classes that help with autocompletion and additional type hints. ## Installation -This package comply to PEP 561. So most of the static code analysers will recognise the stubs automatically. +This package comply to [PEP 561](https://www.python.org/dev/peps/pep-0561/). So most of the static code analysers will recognise the stubs automatically. You just have to type: ``` @@ -15,6 +15,13 @@ pip install construct-typing ## Usage I'm mostly working with VSCode and Pylance (which works really great) ??? But i have also tested the stubs with mypy. ???? +## Tests +The stubs are tested against the pytests of the *construct* package in a slightly modified form. Since the tests are relatively detailed I think most cases are covered. +The new typed constructs have new written tests. + +The tests do not generate errors with: +- mypy (Version TODO) +- pyright (Version TODO) ## Explanation ### Stubs @@ -34,32 +41,67 @@ The problem is to describe the more complex constructs like: Currently only the very unspecific type `typing.Any` can be used as type hint (maybe in the future it can be optimised a little, when [variadic generics](https://mail.python.org/archives/list/typing-sig@python.org/thread/SQVTQYWIOI4TIO7NNBTFFWFMSMS2TA4J/) become available). But the biggest disadvantage is that autocompletion for the named subcons is not available. +Note: The stubs are based on *construct* in Version 2.10. + + ### Typed -To include autocompletion and further enhance the type hints for these complex constructs the **construct_typed** package is used as an extension to the original *construct* package. +TODO: +Es werden die Standard Python Klassen benutzt: +- "dataclasses.dataclass" für Struct, Union, ... (anstatt construct.Container) +- "list" für Array, ... (anstatt construct.ListContainer) +- "enum.Enum" für Enums +- "enum.EnumFlag" für EnumFlags + +TODO: +Es handelt sich in der aktuellen Version noch um einen experimentelle Version! + +TODO: +Es handelt sich um "strongly typed". D.h. es gibt keine Unterscheidung zwischen ParsedType und BuildTypes... Die korrenten Typen werden beim +"build" erzwungen (enforced). Bei einem falschen typen, wird eine exception (TypeError) erzeugt. +Nachteil: dass man manchmal mehr code schreiben muss um den korrekten klassennamen zu deklarieren, anstatt einfach nur "dict" zu schreiben +Vorteil: während der statischen Codeanalyse können schon mehr fehler entdeckt werden. + + +To include autocompletion and further enhance the type hints for these complex constructs the **construct_typed** package is used as an extension to the original *construct* package. It is mainly a bunch of Adapters for the original constructs with the focus on type hints. It implements the following new types: -- TypedEnum -- TypedStruct -- TypedBitStruct -- TypedUnion +- `TStruct`: similar to `construct.Struct` but with `dataclasses.dataclass` +- `TBitStruct`: similar to `construct.BitStruct` but with `dataclasses.dataclass` +- `TEnum`: similar to `construct.Enum` but with `construct_typed.EnumBase` +- `TArray`: similar to `construct.Array` but with `list` insted of `construct.ListContainer` +- TODO: `TUnion` -An example of the added `TypedStruct` class: +A short example: ```python -from construct import Const, Int8ub, Array, this, Byte -from construct_typed import TypedContainer, Subcon, TypedStruct +import dataclasses +import typing as t +import construct as cs +import construct_typed as cst -class Image(TypedContainer): - signature: Subcon(Const(b"BMP")) - width: Subcon(Int8ub()) - height: Subcon(Int8ub()) - pixels: Subcon(Array(cs.this.width * cs.this.height, Byte())) -format = TypedStruct(Image) -obj = Image(width=3, height=2, pixels=[7, 8, 9, 11, 12, 13]) +class Orientation(cst.EnumBase): + HORIZONTAL = 0 + VERTICAL = 1 + +@dataclasses.dataclass +class Image: + signature: t.Optional[bytes] = cst.TStructField(cs.Const(b"BMP")) + orientation: Orientation = cst.TStructField(cst.TEnum(cs.Int8ub, Orientation)) + width: int = cst.TStructField(cs.Int8ub) + height: int = cst.TStructField(cs.Int8ub) + pixels: t.List[int] = cst.TStructField(cst.TArray(cs.this.width * cs.this.height, cs.Byte)) + +format = cst.TStruct(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\x03\x02\x07\x08\t\x0b\x0c\r")) +print(format.parse(b"BMP\x01\x03\x02\x07\x08\t\x0b\x0c\r")) +``` +Output: +``` +b'BMP\x01\x03\x02\x07\x08\t\x0b\x0c\r' +Image(signature=b'BMP', orientation=, width=3, height=2, pixels=[7, 8, 9, 11, 12, 13]) ``` -An example of the added `TypedEnum` class: \ No newline at end of file + diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index c7afc44..a749902 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -7,7 +7,7 @@ from .generic_wrapper import ( PathType, ) from .tarray import TArray -from .tenum import TEnum +from .tenum import TEnum, EnumBase from .tstruct import TBitStruct, TStruct, TStructField from .tunion import TUnion, TUnionField @@ -19,6 +19,7 @@ __all__ = [ "TUnionField", "TUnion", "TArray", + "EnumBase", "Construct", "Adapter", "ListContainer", diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generic_wrapper.py index dead0ba..2f4685f 100644 --- a/construct_typed/generic_wrapper.py +++ b/construct_typed/generic_wrapper.py @@ -39,4 +39,4 @@ else: pass ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]] - PathType = str + PathType = str \ No newline at end of file diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index 8c4439b..c713b46 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -1,73 +1,66 @@ import enum import typing as t -import construct as cs - from .generic_wrapper import * -EnumType = t.TypeVar("EnumType", bound=enum.IntEnum) +class EnumBase(enum.IntEnum): + """ + Base class for an Enum used in `construct_typed.TEnum`. -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 + This class extends the standard `enum.IntEnum`, so that missing values are automatically generated. + """ + # 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 + @classmethod + def _missing_( + cls, value: t.Any + ) -> t.Optional["EnumBase"]: + 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, value: int) -> "EnumBase": + 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 # type: ignore + +EnumType = t.TypeVar("EnumType", bound=EnumBase) + +class TEnum(Adapter[int, int, EnumType, EnumType]): + """ + Typed enum. + """ def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): - if not issubclass(enum_type, enum.IntEnum): + if not issubclass(enum_type, EnumBase): raise TypeError( - 'The class "{}" is not an "enum.IntEnum"'.format(enum_type.__name__) + "'{}' has to be a '{}'".format(repr(enum_type), repr(EnumBase)) ) - @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) - def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType: - return self.enum_type(obj) + def _decode(self, obj: int, context: Context, path: PathType) -> EnumType: + return self.enum_type(obj) # type: ignore def _encode( self, - obj: t.Union[int, str, EnumType], + obj: EnumType, context: "cs.Context", path: "cs.PathType", ) -> int: - try: - # TODO: remove this. only strongly typed enums are allowed... - 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 - ) + if isinstance(obj, self.enum_type): + return int(obj) + raise TypeError("'{}' has to be of type {}".format(repr(obj), repr(self.enum_type))) diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py index 4e8d6a2..6196ed5 100644 --- a/construct_typed/tstruct.py +++ b/construct_typed/tstruct.py @@ -44,7 +44,7 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): ) -> None: if not dataclasses.is_dataclass(dataclass_type): raise TypeError( - "'{}' has to be a 'dataclasses.dataclass'".format(dataclass_type) + "'{}' has to be a 'dataclasses.dataclass'".format(repr(dataclass_type)) ) self.dataclass_type = dataclass_type self.swapped = swapped @@ -105,7 +105,7 @@ 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(obj, self.dataclass_type)) + raise TypeError("'{}' has to be of type {}".format(repr(obj), repr(self.dataclass_type))) class TStruct(_TStruct[ParsedType, ParsedType]): diff --git a/tests/test_typed.py b/tests/test_typed.py index d686375..d242bd7 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -122,7 +122,7 @@ def test_tstruct_anonymus_fields_2() -> None: assert d.build(TestDataclass()) == d.build(TestDataclass()) -def test_tstruct_missing_dataclass() -> None: +def test_tstruct_no_dataclass() -> None: class TestDataclass: a: int = cst.TStructField(cs.Int16ub) b: int = cst.TStructField(cs.Int8ub) @@ -151,43 +151,32 @@ def test_tbitstruct() -> None: def test_tenum() -> None: - class E(enum.IntEnum): + class E(cst.EnumBase): a = 1 b = 2 common(cst.TEnum(cs.Byte, E), b"\x01", E.a, 1) - common(cst.TEnum(cs.Byte, E), b"\x01", 1, 1) + common(cst.TEnum(cs.Byte, E), b"\x02", E.b, 1) + common(cst.TEnum(cs.Byte, E), b"\x03", E(3), 1) + common(cst.TEnum(cs.Byte, E), b"\xff", E(255), 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(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" - - -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" + assert obj == E(3) + assert obj == 3 + obj = format.parse(b"\xff") + assert obj == E(255) + assert obj == 255 -def test_tenum_no_int_enum() -> None: + +def test_tenum_no_enumbase() -> None: class E(enum.Enum): a = 1 b = 2 @@ -195,8 +184,22 @@ def test_tenum_no_int_enum() -> None: assert raises(lambda: cst.TEnum(cs.Byte, E)) == TypeError +def test_tstruct_wrong_enumbase() -> None: + class E1(cst.EnumBase): + a = 1 + b = 2 + + class E2(cst.EnumBase): + a = 1 + b = 2 + + assert ( + raises(cst.TEnum(cs.Byte, E1).build, E2.a) == TypeError + ) + + def test_tenum_in_tstruct() -> None: - class TestEnum(enum.IntEnum): + class TestEnum(cst.EnumBase): a = 1 b = 2 @@ -207,7 +210,11 @@ def test_tenum_in_tstruct() -> None: common( cst.TStruct(TestDataclass), - b"\x00\x01\x02", - TestDataclass(a=TestEnum.a, b=TestEnum.b), - 3, + b"\x01\x02", + TestDataclass(a=TestEnum.a, b=2), + 2, + ) + + assert ( + raises(cst.TEnum(cs.Byte, TestEnum).build, TestDataclass(a=1, b=2)) == TypeError # type: ignore )