Add SubsetEnum, SupersetEnum, and StrictEnum

src/fastcs/attributes/attr_r.py

@@ -46,6 +46,13 @@ def get(self) -> DType_T:
         """Get the cached value of the attribute."""
         return self._value
 
+    def update_datatype(self, datatype: DataType[DType_T]) -> None:
+        value = self._value
+        if not isinstance(value, datatype.dtype):  # e.g. a member of a replaced enum
+            value = datatype.validate(value)
+        super().update_datatype(datatype)
+        self._value = value
+
     @property
     def access_mode(self) -> AttributeAccessMode:
         return "r"

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 import Enum, EnumTypes, StrictEnum, validate_enum_choices
 from fastcs.logging import logger
 from fastcs.methods import Command, Method, Scan, UnboundCommand, UnboundScan
 from fastcs.tracer import Tracer
@@ -323,7 +325,13 @@ 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:
+            if (
+                hint.dtype is not None
+                and isinstance(attr.datatype, Enum)
+                and issubclass(hint.dtype, EnumTypes)
+            ):
+                self._validate_hinted_enum(name, hint.dtype, attr)
+            elif 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. "
@@ -336,6 +344,23 @@ def add_attribute(self, name, attr: Attribute):
         self.__attributes[name] = attr
         super().__setattr__(name, attr)
 
+    def _validate_hinted_enum(
+        self, name: str, hinted: type[enum.Enum], attr: Attribute
+    ) -> None:
+        """Validate introspected enum choices and adopt the hinted enum"""
+        introspected = attr.datatype.dtype
+        try:
+            validate_enum_choices(hinted, introspected)
+        except TypeError as exc:
+            raise RuntimeError(
+                f"Controller '{self.__class__.__name__}' introspection of "
+                f"hinted attribute '{name}' does not match defined enum choices. {exc}"
+            ) from None
+
+        # Superset and Subset serve only the choices the device has
+        allowed = None if issubclass(hinted, StrictEnum) else introspected
+        attr.update_datatype(Enum(hinted, allowed=allowed))
+
     @property
     def attributes(self) -> dict[str, Attribute]:
         return self.__attributes

src/fastcs/datatypes/__init__.py

@@ -4,6 +4,11 @@
 from .datatype import DType as DType
 from .datatype import DType_T as DType_T
 from .enum import Enum as Enum
+from .enum import EnumTypes as EnumTypes
+from .enum import StrictEnum as StrictEnum
+from .enum import SubsetEnum as SubsetEnum
+from .enum import SupersetEnum as SupersetEnum
+from .enum import validate_enum_choices as validate_enum_choices
 from .float import Float as Float
 from .int import Int as Int
 from .string import String as String

src/fastcs/datatypes/enum.py

@@ -1,7 +1,7 @@
 import enum
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from fastcs.datatypes.datatype import DataType
 
@@ -12,6 +12,9 @@
 @dataclass(frozen=True)
 class Enum(Generic[Enum_T], DataType[Enum_T]):
     enum_cls: type[Enum_T]
+    allowed: type[enum.Enum] | None = None
+    """Introspected device enum restricting (for SupersetEnum)
+      or extending (for SubsetEnum) the choices of ``enum_cls``, mapped by value"""
 
     def __post_init__(self):
         if not issubclass(self.enum_cls, enum.Enum):
@@ -22,11 +25,29 @@ def index_of(self, value: Enum_T) -> int:
 
     @cached_property
     def members(self) -> list[Enum_T]:
-        return list(self.enum_cls)
+        if self.allowed is None:
+            return list(self.enum_cls)
+        # SubsetEnum returns the raw value for device choices it doesn't define
+        return [self.enum_cls(device_member.value) for device_member in self.allowed]
 
     @cached_property
     def names(self) -> list[str]:
-        return [member.name for member in self.members]
+        if self.allowed is None:
+            return [member.name for member in self.members]
+        return [
+            choice.name if isinstance(choice, self.enum_cls) else device_member.name
+            for choice, device_member in zip(self.members, self.allowed, strict=True)
+        ]
+
+    def validate(self, value: Any) -> Enum_T:
+        # Map members of another enum, e.g. an introspected one, by value
+        if isinstance(value, enum.Enum) and not isinstance(value, self.enum_cls):
+            value = value.value
+
+        value = super().validate(value)
+        if value not in self.members:
+            raise ValueError(f"{value!r} is not one of the choices {self.names}")
+        return value
 
     @property
     def dtype(self) -> type[Enum_T]:
@@ -35,3 +56,77 @@ def dtype(self) -> type[Enum_T]:
     @property
     def initial_value(self) -> Enum_T:
         return self.members[0]
+
+
+class StrictEnum(enum.Enum):
+    """All members should exist on the device, and there will be no extras."""
+
+
+class _SubsetEnumMeta(enum.EnumMeta):
+    _member_type_: type
+    """The mixed in value type, e.g. ``str``, set by ``EnumMeta``"""
+
+    def __new__(mcls, *args, **kwargs):
+        cls = super().__new__(mcls, *args, **kwargs)
+        if cls._member_map_ and cls._member_type_ is object:
+            raise TypeError(
+                f"{cls.__name__} needs a value type, e.g. (str, SubsetEnum)"
+            )
+        return cls
+
+    def __call__(cls, value, names=None, *args, **kwargs):  # type: ignore
+        """Return the member for ``value``, or ``value`` itself if it is of the
+        member type but not one of the members."""
+        if names is not None:  # functional API
+            return super().__call__(value, names, *args, **kwargs)
+        try:
+            return super().__call__(value)
+        except ValueError:
+            if isinstance(value, cls._member_type_):
+                return value
+            raise
+
+
+class SubsetEnum(enum.Enum, metaclass=_SubsetEnumMeta):
+    """All members should exist on the device, but the device may have extras.
+
+    Subclasses must mix in a value type, e.g. ``class Trigger(str, SubsetEnum)``, as
+    values for device choices outside the enum are kept as that raw type.
+    """
+
+
+class SupersetEnum(enum.Enum):
+    """Some members should exist on the device, and there should be no extras."""
+
+
+EnumTypes = (StrictEnum, SubsetEnum, SupersetEnum)
+
+
+def validate_enum_choices(
+    hinted: type[enum.Enum], introspected: type[enum.Enum]
+) -> None:
+    """Check introspected enum choices against a hinted enum, compared by value.
+
+    Raises:
+        TypeError: If the choices are not compatible with the hinted enum type
+
+    """
+    choices = [member.value for member in hinted]
+    device_choices = [member.value for member in introspected]
+    error_msg = (
+        f"Device has choices {device_choices}, but {hinted.__name__} requested "
+        f"{choices} "
+    )
+    if issubclass(hinted, StrictEnum):
+        if set(choices) != set(device_choices):
+            raise TypeError(error_msg + "to be strictly equal to them.")
+    elif issubclass(hinted, SubsetEnum):
+        if not set(choices).issubset(device_choices):
+            raise TypeError(error_msg + "to be a subset of them.")
+    elif issubclass(hinted, SupersetEnum):
+        if not set(device_choices).issubset(choices):
+            raise TypeError(error_msg + "to be a superset of them.")
+    else:
+        raise TypeError(
+            f"{hinted.__name__} is not a StrictEnum, SubsetEnum, or SupersetEnum"
+        )

src/fastcs/transports/epics/ca/util.py

@@ -1,4 +1,3 @@
-import enum
 import re
 from collections.abc import Callable
 from dataclasses import asdict
@@ -246,10 +245,8 @@ def cast_from_epics_type(datatype: DataType[DType_T], value: object) -> DType_T:
                 return datatype.validate(datatype.members[value])
             else:  # enum backed by string record
                 assert isinstance(value, str), "Got non-string value for long Enum"
-                # python typing can't narrow the nested generic enum_cls
-                assert issubclass(datatype.enum_cls, enum.Enum), "Invalid Enum.enum_cls"
-                enum_member = datatype.enum_cls[value]
-                return datatype.validate(enum_member)
+                # Map by position, as SubsetEnum extras are raw values, not members
+                return datatype.validate(datatype.members[datatype.names.index(value)])
         case datatype if issubclass(type(datatype), EPICS_ALLOWED_DATATYPES):
             return datatype.validate(value)  # type: ignore
         case _:
@@ -263,7 +260,8 @@ def cast_to_epics_type(datatype: DataType[DType_T], value: DType_T) -> Any:
             if len(datatype.members) <= MBB_MAX_CHOICES:
                 return datatype.index_of(datatype.validate(value))
             else:  # enum backed by string record
-                return datatype.validate(value).name
+                # Map by position, as SubsetEnum extras are raw values, not members
+                return datatype.names[datatype.index_of(datatype.validate(value))]
         case String() as string:
             if string.length is not None:
                 return value[: string.length]

src/fastcs/transports/tango/util.py

@@ -2,7 +2,7 @@
 from dataclasses import asdict
 from typing import Any
 
-from tango import AttrDataFormat
+from tango import AttrDataFormat, CmdArgType
 
 from fastcs.attributes import Attribute
 from fastcs.datatypes import (
@@ -100,6 +100,10 @@ def get_server_metadata_from_datatype(datatype: DataType[DType]) -> dict[str, st
                     )
         case Float():
             arguments["format"] = f"%.{datatype.prec}"
+        case Enum():
+            # Labels must match the served choices that indices refer to
+            dtype = CmdArgType.DevEnum
+            arguments["enum_labels"] = datatype.names
 
     arguments["dtype"] = dtype
     for argument, value in arguments.items():

tests/example_p4p_ioc.py

@@ -10,6 +10,7 @@
 from fastcs.launch import FastCS
 from fastcs.methods import command, scan
 from fastcs.transports.epics.pva import EpicsPVATransport
+from tests.util import DeviceGain, DeviceMode, DeviceTrigger, Gain, Mode, Trigger
 
 
 @dataclass
@@ -42,10 +43,22 @@ class ParentController(Controller):
         Table([("A", np.int32), ("B", "i"), ("C", "?"), ("D", np.float64)]),
         io_ref=SimpleAttributeIORef(),
     )
+    mode: AttrRW[Mode]
+    gain: AttrRW[Gain]
+    trigger: AttrRW[Trigger]
 
     def __init__(self, description=None, ios=None):
         super().__init__(description, ios)
 
+    async def initialise(self):
+        # Enums as introspected from a device, validated against the hints above
+        for name, device_enum in (
+            ("mode", DeviceMode),
+            ("gain", DeviceGain),
+            ("trigger", DeviceTrigger),
+        ):
+            self.add_attribute(name, AttrRW(Enum(device_enum)))
+
 
 class ChildController(Controller):
     fail_on_next_e = True

tests/example_softioc.py

@@ -1,21 +1,57 @@
 import asyncio
+import enum
 from pathlib import Path
+from typing import cast
 
 from fastcs.attributes import AttrR, AttrRW, AttrW
 from fastcs.control_system import FastCS
 from fastcs.controllers import Controller, ControllerVector
-from fastcs.datatypes import Int
+from fastcs.datatypes import Enum, Int, SubsetEnum
 from fastcs.methods import command
 from fastcs.transports.epics.ca.transport import (
     EpicsCAOptions,
     EpicsCATransport,
     EpicsGUIOptions,
 )
+from tests.util import DeviceGain, DeviceMode, DeviceTrigger, Gain, Mode, Trigger
+
+
+class LongTrigger(str, SubsetEnum):
+    INTERNAL = "Internal"
+    EXTERNAL = "External"
+
+
+# More than 16 choices, so CA serves it as a string record of names
+DeviceLongTrigger = cast(
+    type[enum.Enum],
+    enum.Enum(
+        "DeviceLongTrigger",
+        {
+            "INTERNAL": "Internal",
+            "EXTERNAL": "External",
+            **{f"LINE{i}": f"Line{i}" for i in range(16)},
+        },
+    ),
+)
 
 
 class ParentController(Controller):
     a: AttrR = AttrR(Int())
     b: AttrRW = AttrRW(Int())
+    mode: AttrRW[Mode]
+    gain: AttrRW[Gain]
+    trigger: AttrRW[Trigger]
+    long_trigger: AttrRW[LongTrigger]
+
+    async def initialise(self):
+        # Enums as introspected from a device, validated against the hints above
+        for name, device_enum in (
+            ("mode", DeviceMode),
+            ("gain", DeviceGain),
+            ("trigger", DeviceTrigger),
+            ("long_trigger", DeviceLongTrigger),
+        ):
+            self.add_attribute(name, AttrRW(Enum(device_enum)))
 
 
 class ChildController(Controller):