diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index 00d5093..e052ee5 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -18,7 +18,7 @@ from .generic_wrapper import ( ListContainer, PathType, ) -from .tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum +from .tenum import EnumBase, EnumValue, FlagsEnumBase, TEnum, TFlagsEnum __all__ = [ "DataclassBitStruct", @@ -32,6 +32,7 @@ __all__ = [ "csfield", "sfield", "EnumBase", + "EnumValue", "FlagsEnumBase", "TEnum", "TFlagsEnum", diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index 6855f3d..3ead0dd 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -5,13 +5,55 @@ from .generic_wrapper import * # ## TEnum ############################################################################################################ +class EnumValue: + """ + This is a helper class for adding documentation to an enum value. + """ + + def __init__(self, value: int, doc: t.Optional[str] = None) -> None: + self.value = value + self.__doc__ = doc if doc else "" + + def __int__(self) -> int: + return self.value + + class EnumBase(enum.IntEnum): """ Base class for an Enum used in `construct_typed.TEnum`. - This class extends the standard `enum.IntEnum`, so that missing values are automatically generated. + This class extends the standard `enum.IntEnum` by. + - missing values are automatically generated + - possibility to add documentation for each enum value (see `EnumValue`) + + Example:: + + >>> class State(EnumBase): + ... Idle = 1 + ... Running = EnumValue(2, "This is the running state.") + + >>> State(1) + + + >>> State["Idle"] + + + >>> State.Idle + + + >>> State(3) # missing value + + + >>> State.Running.__doc__ # documentation + 'This is the running state.' """ + def __init__(self, val: t.Union[EnumValue, int]): + if isinstance(val, EnumValue): + self.__doc__ = val.__doc__ + else: + self.__doc__ = "" + # 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 @@ -25,6 +67,7 @@ class EnumBase(enum.IntEnum): # However, new_member._name_ = value works, too new_member._name_ = str(value) new_member._value_ = value + new_member.__doc__ = "missing value" pseudo_member = cls._value2member_map_.setdefault(value, new_member) return pseudo_member return None # will raise the ValueError in Enum.__new__ @@ -75,7 +118,51 @@ class TEnum(Adapter[int, int, EnumType, EnumType]): # ## TFlagsEnum ####################################################################################################### class FlagsEnumBase(enum.IntFlag): - pass + """ + Base class for an Enum used in `construct_typed.TFlagsEnum`. + + This class extends the standard `enum.IntFlag` by. + - possibility to add documentation for each enum value (see `EnumValue`) + + Example:: + + >>> class Option(FlagsEnumBase): + ... OptOne = 1 + ... OptTwo = EnumValue(2, "This is option two.") + + >>> Option(1) + + + >>> Option["OptOne"] + + + >>> Option.OptOne + + + >>> Option(3) + + + >>> Option(4) + + + >>> Option.OptTwo.__doc__ # documentation + 'This is option two.' + """ + + def __init__(self, val: t.Union[EnumValue, int]): + if isinstance(val, EnumValue): + self.__doc__ = val.__doc__ + else: + self.__doc__ = "" + + @classmethod + def _missing_(cls, value: t.Any) -> t.Any: + """ + Returns member (possibly creating it) if one can be found for value. + """ + new_member = super()._missing_(value) + new_member.__doc__ = "missing value" + return new_member FlagsEnumType = t.TypeVar("FlagsEnumType", bound=FlagsEnumBase) diff --git a/tests/test_typed.py b/tests/test_typed.py index d74e025..756b0f6 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -384,6 +384,42 @@ def test_tenum_no_enumbase() -> None: assert raises(lambda: cst.TEnum(cs.Byte, cls)) == TypeError +def test_tenum_docstring() -> None: + class TestEnum(cst.EnumBase): + """ + This is an test enum. + """ + + Value_WithDoc = cst.EnumValue(0, doc="an enum with a documentation") + Value_WithMultilineDoc = cst.EnumValue( + 1, + """ + An enum with a multiline documentation... + ...next line... + """, + ) + Value_NoDoc = cst.EnumValue(2) + Value_NoDoc2 = 3 + + assert ( + TestEnum.__doc__ + == """ + This is an test enum. + """ + ) + assert TestEnum.Value_WithDoc.__doc__ == "an enum with a documentation" + assert ( + TestEnum.Value_WithMultilineDoc.__doc__ + == """ + An enum with a multiline documentation... + ...next line... + """ + ) + assert TestEnum.Value_NoDoc.__doc__ == "" + assert TestEnum.Value_NoDoc2.__doc__ == "" + assert TestEnum(5).__doc__ == "missing value" + + def test_dataclass_struct_wrong_enumbase() -> None: class E1(cst.EnumBase): a = 1 @@ -434,3 +470,39 @@ def test_tenum_flags() -> None: assert d.build(TestEnum(255)) == b"\xff" assert d.build(TestEnum.eight) == b"\x08" assert raises(d.build, 2) == TypeError + + +def test_tenum_flags_docstring() -> None: + class TestEnum(cst.FlagsEnumBase): + """ + This is an test flags enum. + """ + + Value_WithDoc = cst.EnumValue(0, doc="an enum with a documentation") + Value_WithMultilineDoc = cst.EnumValue( + 1, + """ + An enum with a multiline documentation... + ...next line... + """, + ) + Value_NoDoc = cst.EnumValue(2) + Value_NoDoc2 = 4 + + assert ( + TestEnum.__doc__ + == """ + This is an test flags enum. + """ + ) + assert TestEnum.Value_WithDoc.__doc__ == "an enum with a documentation" + assert ( + TestEnum.Value_WithMultilineDoc.__doc__ + == """ + An enum with a multiline documentation... + ...next line... + """ + ) + assert TestEnum.Value_NoDoc.__doc__ == "" + assert TestEnum.Value_NoDoc2.__doc__ == "" + assert TestEnum(8).__doc__ == "missing value"