Improve Enum Validation

src/fastcs/controllers/base_controller.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import enum
 from collections import Counter
 from collections.abc import Sequence
 from copy import deepcopy
@@ -13,6 +14,7 @@
 
 from fastcs.attributes import AnyAttributeIO, Attribute, AttrR, AttrW, HintedAttribute
 from fastcs.controllers.controller_api import ControllerAPI
+from fastcs.datatypes.enum import check_enum_hint
 from fastcs.logging import logger
 from fastcs.methods import Command, Method, Scan, UnboundCommand, UnboundScan
 from fastcs.tracer import Tracer
@@ -323,19 +325,38 @@ def add_attribute(self, name, attr: Attribute):
                     f"hinted attribute '{name}' does not match defined access mode. "
                     f"Expected '{hint.attr_type.__name__}' got '{type(attr).__name__}'."
                 )
-            if hint.dtype is not None and hint.dtype != attr.datatype.dtype:
-                raise RuntimeError(
-                    f"Controller '{self.__class__.__name__}' introspection of "
-                    f"hinted attribute '{name}' does not match defined datatype. "
-                    f"Expected '{hint.dtype.__name__}', "
-                    f"got '{attr.datatype.dtype.__name__}'."
-                )
+            if hint.dtype is not None:
+                self._check_hinted_dtype(name, hint.dtype, attr.datatype.dtype)
 
         attr.set_name(name)
         attr.set_path(self.path)
         self.__attributes[name] = attr
         super().__setattr__(name, attr)
 
+    def _check_hinted_dtype(self, name: str, hinted: type, actual: type):
+        """Check the dtype of an `Attribute` against the dtype in its type hint.
+
+        Dtypes must be identical, except for enums, which may also match on their
+        values - see `check_enum_hint`.
+        """
+        if hinted is actual:
+            return
+
+        message = (
+            f"Controller '{self.__class__.__name__}' introspection of "
+            f"hinted attribute '{name}' does not match defined datatype. "
+            f"Expected '{hinted.__name__}', got '{actual.__name__}'"
+        )
+
+        if issubclass(hinted, enum.Enum) and issubclass(actual, enum.Enum):
+            try:
+                check_enum_hint(hinted, actual)
+                return
+            except ValueError as e:
+                raise RuntimeError(f"{message}: {e}.") from e
+
+        raise RuntimeError(f"{message}.")
+
     @property
     def attributes(self) -> dict[str, Attribute]:
         return self.__attributes

src/fastcs/datatypes/enum.py

@@ -35,3 +35,81 @@ def dtype(self) -> type[Enum_T]:
     @property
     def initial_value(self) -> Enum_T:
         return self.members[0]
+
+
+def _value_mixin(enum_cls: type[enum.Enum]) -> type | None:
+    """The builtin type an enum's members compare as.
+
+    Returns ``str`` if ``enum_cls`` is a StrEnum, or ``int`` if it is an IntEnum
+    """
+    for mixin in (str, int):
+        if issubclass(enum_cls, mixin):
+            return mixin
+
+    return None
+
+
+def check_enum_hint(hinted: type[enum.Enum], actual: type[enum.Enum]) -> None:
+    """Check that the enum of an `Attribute` satisfies the enum in its type hint.
+
+    The hint is satisfied if:
+
+    - It is the same class (e.g., using an enum registry for introspection to pull from)
+    - It is a member-less enum base, such as `enum.StrEnum`, and ``actual``
+      derives from it (e.g., ``actual`` is a `enum.StrEnum` with members)
+    - Every value in ``hinted`` is also a value of ``actual``, and both share the same
+      ``str`` or ``int`` mixin.
+
+    Args:
+        hinted: The enum class in the type hint
+        actual: The enum class the `Attribute` holds
+
+    Raises:
+        ValueError: Describing why the hint is not satisfied
+
+    """
+    if hinted is actual:  # The same enum class is used
+        return
+
+    if len(hinted.__members__) == 0:  # This is a member-less enum
+        if issubclass(actual, hinted):  # Actual enum derived from the member-less enum
+            return
+        raise ValueError(
+            f"'{actual.__name__}' is not a '{hinted.__name__}'"
+        )  # Raise as different kinds of enums
+
+    hinted_mixin = _value_mixin(hinted)
+    if hinted_mixin is None:
+        # In this case, we have a defined enum.Enum subclass, but hinted is not actual,
+        # so it can only match itself.
+        raise ValueError(
+            f"'{hinted.__name__}' has members but no str or int mixin, so it only "
+            "matches itself - use enum.StrEnum or enum.IntEnum to match an "
+            "introspected enum by value"
+        )
+
+    actual_mixin = _value_mixin(actual)
+    if actual_mixin is not hinted_mixin:
+        actual_kind = actual_mixin.__name__ if actual_mixin else "no"
+        # In this case, we have a defined enum.StrEnum or enum.IntEnum subclass,
+        # but the actual enum is of a different enum type.
+        raise ValueError(
+            f"'{hinted.__name__}' has {hinted_mixin.__name__} values but "
+            f"'{actual.__name__}' has {actual_kind} values"
+        )
+
+    # `__members__` rather than iterating the class, which skips aliases and, for
+    # flags, zero and multi-bit values - all of which are still named values
+    actual_values = {member.value for member in actual.__members__.values()}
+    missing = [
+        member.value
+        for member in hinted.__members__.values()
+        if member.value not in actual_values
+    ]
+    # A hinted enum_cls can have a subset of the members the actual enum_cls has,
+    # but must not contain any extra members.
+    if missing:
+        raise ValueError(
+            f"'{hinted.__name__}' has values that '{actual.__name__}' does not: "
+            f"{', '.join(repr(value) for value in missing)}"
+        )

tests/test_controllers.py

@@ -182,8 +182,8 @@ def test_enum_attribute_hint_validation():
     class GoodEnum(enum.IntEnum):
         VAL = 0
 
-    class BadEnum(enum.IntEnum):
-        VAL = 0
+    class BadEnum(enum.StrEnum):
+        VAL = "val"
 
     class HintedController(Controller):
         enum: AttrRW[GoodEnum]

tests/test_datatypes.py

@@ -1,10 +1,30 @@
+import enum
 from enum import IntEnum
 
 import numpy as np
 import pytest
 
 from fastcs.datatypes import Bool, DataType, Enum, Float, Int, String, Table, Waveform
 from fastcs.datatypes._util import numpy_to_fastcs_datatype
+from fastcs.datatypes.enum import check_enum_hint
+
+
+class State(enum.StrEnum):
+    IDLE = "idle"
+    READY = "ready"
+
+
+class PlainState(enum.Enum):
+    IDLE = "idle"
+
+
+class IntState(enum.IntEnum):
+    IDLE = 0
+
+
+Introspected = enum.StrEnum(
+    "Introspected", {"idle": "idle", "ready": "ready", "error": "error"}
+)
 
 
 def test_base_validate():
@@ -129,3 +149,56 @@ def test_string_length():
 
     with pytest.raises(ValueError):
         String(length=0)
+
+
+@pytest.mark.parametrize(
+    "hinted, actual",
+    [
+        pytest.param(State, State, id="same-class"),
+        pytest.param(enum.StrEnum, Introspected, id="member-less-base"),
+        pytest.param(enum.Enum, IntState, id="enum-base-accepts-any-enum"),
+        pytest.param(State, Introspected, id="values-are-a-subset"),
+    ],
+)
+def test_check_enum_hint_accepts(hinted, actual):
+    check_enum_hint(hinted, actual)  # Does not raise
+
+
+def test_check_enum_hint_rejects_enum_not_derived_from_member_less_hint():
+    with pytest.raises(ValueError, match="'IntState' is not a 'StrEnum'"):
+        check_enum_hint(enum.StrEnum, IntState)
+
+
+def test_check_enum_hint_rejects_hint_with_members_and_no_mixin():
+    with pytest.raises(ValueError, match="'PlainState' has members but no str or int"):
+        check_enum_hint(PlainState, Introspected)
+
+
+def test_check_enum_hint_rejects_mismatched_mixin():
+    with pytest.raises(
+        ValueError, match="'State' has str values but 'IntState' has int values"
+    ):
+        check_enum_hint(State, IntState)
+
+
+def test_check_enum_hint_rejects_hinted_value_missing_from_actual():
+    actual = enum.StrEnum("Actual", {"idle": "idle"})
+
+    with pytest.raises(
+        ValueError, match="'State' has values that 'Actual' does not: 'ready'"
+    ):
+        check_enum_hint(State, actual)
+
+
+def test_check_enum_hint_rejects_named_flag_value_missing_from_actual():
+    class Hinted(enum.IntFlag):
+        A = 1
+        BOTH = 3  # A multi-bit value, which iterating the class skips
+
+    class Actual(enum.IntFlag, boundary=enum.STRICT):
+        A = 1
+
+    with pytest.raises(
+        ValueError, match="'Hinted' has values that 'Actual' does not: 3$"
+    ):
+        check_enum_hint(Hinted, Actual)