Repository: pcdshub/pcdsutils

Issue:
Add HelpfulIntEnum from pcdsdaq
https://github.com/pcdshub/pcdsdaq/pull/108/files#r785021109
https://github.com/pcdshub/pcdsdaq/pull/108/files#r785022116

pcdsdaq/pcdsdaq/daq/interface.py

Comments:

Patch:
diff --git a/pcdsutils/__init__.py b/pcdsutils/__init__.py
index 14ae388..f4f9e12 100644
--- a/pcdsutils/__init__.py
+++ b/pcdsutils/__init__.py
@@ -1,8 +1,8 @@
+from . import log, utils
 from ._version import get_versions
-from . import utils
-from . import log
+from .enum import HelpfulIntEnum
 
 __version__ = get_versions()['version']
 del get_versions
 
-__all__ = ['utils', 'log']
+__all__ = ['utils', 'log', "HelpfulIntEnum"]
diff --git a/pcdsutils/enum.py b/pcdsutils/enum.py
new file mode 100644
index 0000000..2660ab5
--- /dev/null
+++ b/pcdsutils/enum.py
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import enum
+from typing import Iterator, Set, Union
+
+EnumId = Union[enum.Enum, int, str]
+
+
+class CaseInsensitiveEnumMeta(enum.EnumMeta):
+    def __getattr__(self, key: EnumId) -> enum.Enum:
+        if hasattr(key, "lower"):
+            for item in self:
+                if item.name.lower() == key.lower():
+                    return item
+        return super().__getattr__(key)
+
+    def __getitem__(self, key: EnumId) -> enum.Enum:
+        if hasattr(key, "lower"):
+            for item in self:
+                if item.name.lower() == key.lower():
+                    return item
+        return super().__getitem__(key)
+
+
+class HelpfulIntEnum(enum.IntEnum, metaclass=CaseInsensitiveEnumMeta):
+    """
+    IntEnum subclass with some utility extensions and case insensitivity.
+    """
+
+    @classmethod
+    def from_any(cls, identifier: EnumId) -> HelpfulIntEnum:
+        """
+        Try all the ways to interpret identifier as the enum.
+
+        This is intended to consolidate the try/except tree typically used
+        to interpret external input as an enum.
+
+        Parameters
+        ----------
+        identifier : EnumId
+            Any str, int, or Enum value that corresponds with a valid value
+            on this HelpfulIntEnum instance.
+
+        Returns
+        -------
+        enum : HelpfulIntEnum
+            The corresponding enum object associated with the identifier.
+
+        Raises
+        ------
+        ValueError
+            If the value is not a valid enum identifier.
+        """
+        try:
+            return cls[identifier]
+        except KeyError:
+            return cls(identifier)
+
+    @classmethod
+    def include(
+        cls,
+        identifiers: Iterator[EnumId],
+    ) -> Set[HelpfulIntEnum]:
+        """
+        Returns all enum values matching the identifiers given.
+        This is a shortcut for calling cls.from_any many times and
+        assembling a set of the results.
+
+        Parameters
+        ----------
+        identifiers : Iterator[EnumId]
+            Any iterable that contains strings, ints, and Enum values that
+            correspond with valid values on this HelpfulIntEnum instance.
+
+        Returns
+        -------
+        enums : Set[HelpfulIntEnum]
+            A set whose elements are the enum objects associated with the
+            input identifiers.
+        """
+        return {cls.from_any(ident) for ident in identifiers}
+
+    @classmethod
+    def exclude(
+        cls,
+        identifiers: Iterator[EnumId],
+    ) -> Set[HelpfulIntEnum]:
+        """
+        Return all enum values other than the ones given.
+
+        Parameters
+        ----------
+        identifiers : Iterator[EnumId]
+            Any iterable that contains strings, ints, and Enum values that
+            correspond with valid values on this HelpfulIntEnum instance.
+
+        Returns
+        -------
+        enums : Set[HelpfulIntEnum]
+            A set whose elements are the valid enum objects not associated
+            with the input identifiers.
+        """
+        return set(cls.__members__.values()) - cls.include(identifiers)


Test patch:
diff --git a/pcdsutils/tests/test_enum.py b/pcdsutils/tests/test_enum.py
new file mode 100644
index 0000000..2335813
--- /dev/null
+++ b/pcdsutils/tests/test_enum.py
@@ -0,0 +1,79 @@
+import pytest
+
+from ..enum import HelpfulIntEnum
+
+
+def test_subclass():
+    class MyEnum(HelpfulIntEnum):
+        A = 4
+        B = 5
+        C = 6
+
+    assert MyEnum(4) == MyEnum.A
+    assert MyEnum(5) == MyEnum.B
+    assert MyEnum(6) == MyEnum.C
+
+    with pytest.raises(ValueError):
+        MyEnum(7)
+
+    assert MyEnum["A"] == MyEnum.A == 4
+    assert MyEnum["a"] == MyEnum.A == 4
+    assert MyEnum["a"] == MyEnum.a == 4
+    assert MyEnum.from_any("a") == MyEnum.A == 4
+    assert MyEnum.from_any(4) == MyEnum.A == 4
+
+    assert MyEnum["C"] == MyEnum.C == 6
+    assert MyEnum["c"] == MyEnum.C == 6
+    assert MyEnum["c"] == MyEnum.c == 6
+    assert MyEnum.from_any("c") == MyEnum.C == 6
+    assert MyEnum.from_any(6) == MyEnum.C == 6
+
+
+def test_functional():
+    MyEnum = HelpfulIntEnum("MyEnum", {"a": 1, "B": 2}, module=__name__)
+
+    assert MyEnum["A"] == MyEnum.A == 1
+    assert MyEnum["a"] == MyEnum.A == 1
+    assert MyEnum["a"] == MyEnum.a == 1
+    assert MyEnum.from_any("a") == MyEnum.A
+    assert MyEnum.from_any(1) == MyEnum.A
+
+    assert MyEnum["B"] == MyEnum.B == 2
+    assert MyEnum["b"] == MyEnum.B == 2
+    assert MyEnum["b"] == MyEnum.b == 2
+    assert MyEnum.from_any("b") == MyEnum.B == 2
+    assert MyEnum.from_any(2) == MyEnum.B == 2
+
+
+def test_functional_list():
+    MyEnum = HelpfulIntEnum("MyEnum", ["a", "b"], start=1, module=__name__)
+
+    assert MyEnum["A"] == MyEnum.A == 1
+    assert MyEnum["a"] == MyEnum.A == 1
+    assert MyEnum["a"] == MyEnum.a == 1
+    assert MyEnum.from_any("a") == MyEnum.A
+    assert MyEnum.from_any(1) == MyEnum.A
+
+    assert MyEnum["B"] == MyEnum.B == 2
+    assert MyEnum["b"] == MyEnum.B == 2
+    assert MyEnum["b"] == MyEnum.b == 2
+    assert MyEnum.from_any("b") == MyEnum.B == 2
+    assert MyEnum.from_any(2) == MyEnum.B == 2
+
+
+def test_include():
+    class MyEnum(HelpfulIntEnum):
+        A = 4
+        B = 5
+        C = 6
+
+    assert MyEnum.include([4, "c"]) == {MyEnum.A, MyEnum.C}
+
+
+def test_exclude():
+    class MyEnum(HelpfulIntEnum):
+        A = 4
+        B = 5
+        C = 6
+
+    assert MyEnum.exclude([4, "c"]) == {MyEnum.B}