enhanced EnumBase and FlagsEnumBase to support induvidual documentation for each enum value via EnumValue

Author: timrid

construct_typed/__init__.py

@@ -18,7 +18,7 @@
     ListContainer,
     PathType,
 )
-from .tenum import EnumBase, FlagsEnumBase, TEnum, TFlagsEnum
+from .tenum import EnumBase, EnumValue, FlagsEnumBase, TEnum, TFlagsEnum
 
 __all__ = [
     "DataclassBitStruct",
@@ -32,6 +32,7 @@
     "csfield",
     "sfield",
     "EnumBase",
+    "EnumValue",
     "FlagsEnumBase",
     "TEnum",
     "TFlagsEnum",

construct_typed/tenum.py

@@ -5,13 +5,55 @@
 
 
 # ## 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: 1>
+
+        >>> State["Idle"]
+        <State.Idle: 1>
+
+        >>> State.Idle
+        <State.Idle: 1>
+
+        >>> State(3)  # missing value
+        <State.3: 3>
+
+        >>> 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 @@ def _missing_(cls, value: t.Any) -> t.Optional[enum.Enum]:
                 # 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 @@ def _encode(
 
 # ## 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: 1>
+
+        >>> Option["OptOne"]
+        <Option.OptOne: 1>
+
+        >>> Option.OptOne
+        <Option.OptOne: 1>
+
+        >>> Option(3)
+        <Option.OptTwo|OptOne: 3>
+
+        >>> Option(4)
+        <Option.4: 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)

tests/test_typed.py

@@ -384,6 +384,42 @@ class E(enum.Enum):
     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 @@ class TestEnum(cst.FlagsEnumBase):
     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"