diff --git a/src/fastcs/controllers/base_controller.py b/src/fastcs/controllers/base_controller.py index 8b29ed6d2..df42fb22a 100755 --- a/src/fastcs/controllers/base_controller.py +++ b/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 diff --git a/src/fastcs/datatypes/enum.py b/src/fastcs/datatypes/enum.py index e490f5c76..361100a66 100644 --- a/src/fastcs/datatypes/enum.py +++ b/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)}" + ) diff --git a/tests/test_controllers.py b/tests/test_controllers.py index 5d5dfb4ed..b0fda96b2 100644 --- a/tests/test_controllers.py +++ b/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] diff --git a/tests/test_datatypes.py b/tests/test_datatypes.py index 221fa1de6..3b8c87862 100644 --- a/tests/test_datatypes.py +++ b/tests/test_datatypes.py @@ -1,3 +1,4 @@ +import enum from enum import IntEnum import numpy as np @@ -5,6 +6,25 @@ 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)