From 8169f0ed318a497f5d9c884d153ced16e056ed9d Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Sun, 13 Feb 2022 16:51:20 +0100 Subject: [PATCH] Changed implementation of TEnum and TEnumFlags. --- construct_typed/__init__.py | 22 ++------ construct_typed/tenum.py | 101 ++++++++++++++++++++++++++++-------- 2 files changed, 84 insertions(+), 39 deletions(-) diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index a11e26d..10e5b12 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -1,14 +1,7 @@ from construct_typed.dataclass_struct import ( DataclassBitStruct, - DataclassMixin, DataclassStruct, - TBitStruct, - TContainerBase, - TContainerMixin, - TStruct, - TStructField, csfield, - sfield, ) from construct_typed.generic import ( Adapter, @@ -18,23 +11,16 @@ from construct_typed.generic import ( ListContainer, PathType, ) -from construct_typed.tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum +from construct_typed.tenum import TEnum, TFlags, TEnumConstruct, TFlagsConstruct __all__ = [ "DataclassBitStruct", - "DataclassMixin", "DataclassStruct", - "TBitStruct", - "TContainerBase", - "TContainerMixin", - "TStruct", - "TStructField", "csfield", - "sfield", - "EnumBase", - "FlagsEnumBase", "TEnum", - "TFlagsEnum", + "TEnumConstruct", + "TFlags", + "TFlagsConstruct", "Adapter", "ConstantOrContextLambda", "Construct", diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index 91c4173..4a84584 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -1,28 +1,54 @@ import enum import typing as t -from construct_typed.generic import * +import construct as cs + +from .generic import * + +T = t.TypeVar("T") # ## TEnum ############################################################################################################ -class EnumBase(enum.IntEnum): +class TEnum(enum.IntEnum): """ - Base class for an Enum used in `construct_typed.TEnum`. + Base class for an Enum used in `construct_typed.TEnumConstruct`. 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 = TEnumConstruct(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 @classmethod - def _missing_(cls, value: t.Any) -> t.Optional["EnumBase"]: + def _missing_(cls, value: t.Any) -> t.Optional["TEnum"]: 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": + def _create_pseudo_member_(cls, value: int) -> "TEnum": pseudo_member = cls._value2member_map_.get(value, None) # type: ignore if pseudo_member is None: new_member = int.__new__(cls, value) @@ -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: -EnumType = t.TypeVar("EnumType", bound=EnumBase) + @classmethod + def __construct__(cls: "t.Type[EnumType]") -> "TEnumConstruct[EnumType]": + ... -class TEnum(Adapter[int, int, EnumType, EnumType]): +EnumType = t.TypeVar("EnumType", bound=TEnum) + + +class TEnumConstruct(Adapter[int, int, EnumType, EnumType]): """ Typed enum. """ @@ -46,20 +78,20 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): def __new__( cls, subcon: Construct[int, int], enum_type: t.Type[EnumType] - ) -> "TEnum[EnumType]": + ) -> "TEnumConstruct[EnumType]": ... def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): - if not issubclass(enum_type, EnumBase): + if not issubclass(enum_type, TEnum): raise TypeError( - "'{}' has to be a '{}'".format(repr(enum_type), repr(EnumBase)) + "'{}' has to be a '{}'".format(repr(enum_type), repr(TEnum)) ) # save enum type self.enum_type = t.cast(t.Type[EnumType], enum_type) # type: ignore # init adatper - super(TEnum, self).__init__(subcon) # type: ignore + super(TEnumConstruct, self).__init__(subcon) # type: ignore def _decode(self, obj: int, context: Context, path: PathType) -> EnumType: return self.enum_type(obj) @@ -76,38 +108,65 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): "'{}' has to be of type {}".format(repr(obj), repr(self.enum_type)) ) +# ## TFlags ####################################################################################################### +class TFlags(enum.IntFlag): + @classmethod + def __init_subclass__( + cls, + subcon: "cs.Construct[t.Any, t.Any]", + **kwargs: t.Any, + ): + super().__init_subclass__(**kwargs) -# ## TFlagsEnum ####################################################################################################### -class FlagsEnumBase(enum.IntFlag): - pass + # 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 = TFlagsConstruct(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]", + ) -> "TFlagsConstruct[FlagsEnumType]": + ... -FlagsEnumType = t.TypeVar("FlagsEnumType", bound=FlagsEnumBase) +FlagsEnumType = t.TypeVar("FlagsEnumType", bound=TFlags) -class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): +class TFlagsConstruct(Adapter[int, int, FlagsEnumType, FlagsEnumType]): """ - Typed enum. + Typed flags. """ if t.TYPE_CHECKING: def __new__( cls, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType] - ) -> "TFlagsEnum[FlagsEnumType]": + ) -> "TFlagsConstruct[FlagsEnumType]": ... def __init__(self, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType]): - if not issubclass(enum_type, FlagsEnumBase): + if not issubclass(enum_type, TFlags): raise TypeError( - "'{}' has to be a '{}'".format(repr(enum_type), repr(FlagsEnumBase)) + "'{}' has to be a '{}'".format(repr(enum_type), repr(TFlags)) ) # save enum type self.enum_type = t.cast(t.Type[FlagsEnumType], enum_type) # type: ignore # init adatper - super(TFlagsEnum, self).__init__(subcon) # type: ignore + super(TFlagsConstruct, self).__init__(subcon) # type: ignore def _decode(self, obj: int, context: Context, path: PathType) -> FlagsEnumType: return self.enum_type(obj)