diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index e090f95..e02a688 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -16,9 +16,7 @@ class EnumBase(enum.IntEnum): # 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["EnumBase"]: if isinstance(value, int): return cls._create_pseudo_member_(value) return None # will raise the ValueError in Enum.__new__ @@ -27,7 +25,7 @@ class EnumBase(enum.IntEnum): def _create_pseudo_member_(cls, value: int) -> "EnumBase": pseudo_member = cls._value2member_map_.get(value, None) # type: ignore if pseudo_member is None: - new_member = int.__new__(cls, value) + 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) @@ -35,12 +33,22 @@ class EnumBase(enum.IntEnum): pseudo_member = cls._value2member_map_.setdefault(value, new_member) # type: ignore return pseudo_member # type: ignore + EnumType = t.TypeVar("EnumType", bound=EnumBase) + class TEnum(Adapter[int, int, EnumType, EnumType]): """ Typed enum. """ + + if t.TYPE_CHECKING: + + def __new__( + cls, subcon: Construct[int, int], enum_type: t.Type[EnumType] + ) -> Adapter[int, int, EnumType, EnumType]: + ... + def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]): if not issubclass(enum_type, EnumBase): raise TypeError( @@ -48,7 +56,7 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): ) # save enum type - self.enum_type = t.cast(t.Type[EnumType], enum_type) # type: ignore + self.enum_type = t.cast(t.Type[EnumType], enum_type) # type: ignore # init adatper super(TEnum, self).__init__(subcon) # type: ignore @@ -64,8 +72,9 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): ) -> int: if isinstance(obj, self.enum_type): return int(obj) - raise TypeError("'{}' has to be of type {}".format(repr(obj), repr(self.enum_type))) - + raise TypeError( + "'{}' has to be of type {}".format(repr(obj), repr(self.enum_type)) + ) # ## TFlagsEnum ####################################################################################################### @@ -75,10 +84,19 @@ class FlagsEnumBase(enum.IntFlag): FlagsEnumType = t.TypeVar("FlagsEnumType", bound=FlagsEnumBase) + class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): """ Typed enum. """ + + if t.TYPE_CHECKING: + + def __new__( + cls, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType] + ) -> Adapter[int, int, FlagsEnumType, FlagsEnumType]: + ... + def __init__(self, subcon: Construct[int, int], enum_type: t.Type[FlagsEnumType]): if not issubclass(enum_type, FlagsEnumBase): raise TypeError( @@ -86,7 +104,7 @@ class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): ) # save enum type - self.enum_type = t.cast(t.Type[FlagsEnumType], enum_type) # type: ignore + self.enum_type = t.cast(t.Type[FlagsEnumType], enum_type) # type: ignore # init adatper super(TFlagsEnum, self).__init__(subcon) # type: ignore @@ -102,4 +120,6 @@ class TFlagsEnum(Adapter[int, int, FlagsEnumType, FlagsEnumType]): ) -> int: if isinstance(obj, self.enum_type): return int(obj) - raise TypeError("'{}' has to be of type {}".format(repr(obj), repr(self.enum_type))) + raise TypeError( + "'{}' has to be of type {}".format(repr(obj), repr(self.enum_type)) + )