From 9165c5ae811074b92a0d549946fe85915c538c7a Mon Sep 17 00:00:00 2001 From: Eva Lott Date: Thu, 20 Jun 2024 13:05:32 +0100 Subject: [PATCH] added runtime enum class, metaclass, and tests (#341) added `SubsetEnum` class (with metaclass), and added logic to `make_converter`s to support them Also made it so that * if the record is mbb then the signal has to be `Enum`,`SubsetEnum`, or None * `SubsetEnum` metaclass caches with a tuple, so order matters and `SubsetEnum['A', 'B'] is not `SubsetEnum['B', 'A']` --- src/ophyd_async/core/__init__.py | 91 +++++++++++---------- src/ophyd_async/core/signal_backend.py | 48 ++++++++++- src/ophyd_async/core/soft_signal_backend.py | 41 +++++++--- src/ophyd_async/epics/_backend/common.py | 34 +++++--- tests/core/test_soft_signal_backend.py | 2 +- tests/core/test_subset_enum.py | 90 ++++++++++++++++++++ tests/epics/demo/test_demo.py | 2 +- tests/epics/test_signals.py | 34 ++++++-- 8 files changed, 266 insertions(+), 76 deletions(-) create mode 100644 tests/core/test_subset_enum.py diff --git a/src/ophyd_async/core/__init__.py b/src/ophyd_async/core/__init__.py index a5351757ec..f4f8459cb1 100644 --- a/src/ophyd_async/core/__init__.py +++ b/src/ophyd_async/core/__init__.py @@ -50,7 +50,7 @@ soft_signal_rw, wait_for_value, ) -from .signal_backend import SignalBackend +from .signal_backend import RuntimeSubsetEnum, SignalBackend, SubsetEnum from .soft_signal_backend import SoftSignalBackend from .standard_readable import ConfigSignal, HintedSignal, StandardReadable from .utils import ( @@ -68,66 +68,69 @@ ) __all__ = [ - "get_mock_put", - "callback_on_mock_put", - "mock_puts_blocked", - "set_mock_values", - "reset_mock_put_calls", - "SignalBackend", - "SoftSignalBackend", + "AsyncStatus", + "CalculatableTimeout", + "CalculateTimeout", + "Callback", + "ConfigSignal", + "DEFAULT_TIMEOUT", "DetectorControl", - "MockSignalBackend", "DetectorTrigger", "DetectorWriter", - "StandardDetector", "Device", "DeviceCollector", "DeviceVector", - "Signal", - "SignalR", - "SignalW", - "SignalRW", - "SignalX", - "soft_signal_r_and_setter", - "soft_signal_rw", - "observe_value", - "set_and_wait_for_value", - "set_mock_put_proceeds", - "set_mock_value", - "wait_for_value", - "AsyncStatus", - "WatchableAsyncStatus", "DirectoryInfo", "DirectoryProvider", + "HardwareTriggeredFlyable", + "HintedSignal", + "MockSignalBackend", "NameProvider", + "NotConnected", + "ReadingValueCallback", + "RuntimeSubsetEnum", + "SubsetEnum", "ShapeProvider", - "StaticDirectoryProvider", + "Signal", + "SignalBackend", + "SignalR", + "SignalRW", + "SignalW", + "SignalX", + "SoftSignalBackend", + "StandardDetector", "StandardReadable", - "ConfigSignal", - "HintedSignal", + "StaticDirectoryProvider", + "T", "TriggerInfo", "TriggerLogic", - "HardwareTriggeredFlyable", - "CalculateTimeout", - "CalculatableTimeout", - "DEFAULT_TIMEOUT", - "Callback", - "NotConnected", - "ReadingValueCallback", - "T", + "WatchableAsyncStatus", + "assert_configuration", + "assert_emitted", + "assert_mock_put_called_with", + "assert_reading", + "assert_value", + "callback_on_mock_put", "get_dtype", - "get_unique", - "merge_gathered_dicts", - "wait_for_connection", + "get_mock_put", "get_signal_values", + "get_unique", + "load_device", "load_from_yaml", + "merge_gathered_dicts", + "mock_puts_blocked", + "observe_value", + "reset_mock_put_calls", + "save_device", "save_to_yaml", + "set_and_wait_for_value", + "set_mock_put_proceeds", + "set_mock_value", + "set_mock_values", "set_signal_values", + "soft_signal_r_and_setter", + "soft_signal_rw", + "wait_for_connection", + "wait_for_value", "walk_rw_signals", - "load_device", - "save_device", - "assert_reading", - "assert_value", - "assert_configuration", - "assert_emitted", ] diff --git a/src/ophyd_async/core/signal_backend.py b/src/ophyd_async/core/signal_backend.py index ef8e983421..10498e83b0 100644 --- a/src/ophyd_async/core/signal_backend.py +++ b/src/ophyd_async/core/signal_backend.py @@ -1,5 +1,13 @@ from abc import abstractmethod -from typing import Generic, Optional, Type +from typing import ( + TYPE_CHECKING, + ClassVar, + Generic, + Literal, + Optional, + Tuple, + Type, +) from bluesky.protocols import DataKey, Reading @@ -45,3 +53,41 @@ async def get_setpoint(self) -> T: @abstractmethod def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None: """Observe changes to the current value, timestamp and severity""" + + +class _RuntimeSubsetEnumMeta(type): + def __str__(cls): + if hasattr(cls, "choices"): + return f"SubsetEnum{list(cls.choices)}" + return "SubsetEnum" + + def __getitem__(cls, _choices): + if isinstance(_choices, str): + _choices = (_choices,) + else: + if not isinstance(_choices, tuple) or not all( + isinstance(c, str) for c in _choices + ): + raise TypeError( + "Choices must be a str or a tuple of str, " f"not {type(_choices)}." + ) + if len(set(_choices)) != len(_choices): + raise TypeError("Duplicate elements in runtime enum choices.") + + class _RuntimeSubsetEnum(cls): + choices = _choices + + return _RuntimeSubsetEnum + + +class RuntimeSubsetEnum(metaclass=_RuntimeSubsetEnumMeta): + choices: ClassVar[Tuple[str, ...]] + + def __init__(self): + raise RuntimeError("SubsetEnum cannot be instantiated") + + +if TYPE_CHECKING: + SubsetEnum = Literal +else: + SubsetEnum = RuntimeSubsetEnum diff --git a/src/ophyd_async/core/soft_signal_backend.py b/src/ophyd_async/core/soft_signal_backend.py index bfd4fdbf41..415c167089 100644 --- a/src/ophyd_async/core/soft_signal_backend.py +++ b/src/ophyd_async/core/soft_signal_backend.py @@ -3,14 +3,23 @@ import inspect import time from collections import abc -from dataclasses import dataclass from enum import Enum -from typing import Dict, Generic, Optional, Type, TypedDict, Union, cast, get_origin +from typing import ( + Dict, + Generic, + Optional, + Tuple, + Type, + TypedDict, + Union, + cast, + get_origin, +) import numpy as np from bluesky.protocols import DataKey, Dtype, Reading -from .signal_backend import SignalBackend +from .signal_backend import RuntimeSubsetEnum, SignalBackend from .utils import DEFAULT_TIMEOUT, ReadingValueCallback, T, get_dtype primitive_dtypes: Dict[type, Dtype] = { @@ -74,23 +83,27 @@ def make_initial_value(self, datatype: Optional[Type[T]]) -> T: return cast(T, datatype(shape=0)) # type: ignore -@dataclass class SoftEnumConverter(SoftConverter): - enum_class: Type[Enum] + choices: Tuple[str, ...] - def write_value(self, value: Union[Enum, str]) -> Enum: + def __init__(self, datatype: Union[RuntimeSubsetEnum, Enum]): + if issubclass(datatype, Enum): + self.choices = tuple(v.value for v in datatype) + else: + self.choices = datatype.choices + + def write_value(self, value: Union[Enum, str]) -> str: if isinstance(value, Enum): + return value.value + else: # Runtime enum return value - else: - return self.enum_class(value) def get_datakey(self, source: str, value, **metadata) -> DataKey: - choices = [e.value for e in self.enum_class] return { "source": source, "dtype": "string", "shape": [], - "choices": choices, + "choices": self.choices, **metadata, } @@ -98,13 +111,17 @@ def make_initial_value(self, datatype: Optional[Type[T]]) -> T: if datatype is None: return cast(T, None) - return cast(T, list(datatype.__members__.values())[0]) # type: ignore + if issubclass(datatype, Enum): + return cast(T, list(datatype.__members__.values())[0]) # type: ignore + return cast(T, self.choices[0]) def make_converter(datatype): is_array = get_dtype(datatype) is not None is_sequence = get_origin(datatype) == abc.Sequence - is_enum = issubclass(datatype, Enum) if inspect.isclass(datatype) else False + is_enum = inspect.isclass(datatype) and ( + issubclass(datatype, Enum) or issubclass(datatype, RuntimeSubsetEnum) + ) if is_array or is_sequence: return SoftArrayConverter() diff --git a/src/ophyd_async/epics/_backend/common.py b/src/ophyd_async/epics/_backend/common.py index e5a7453823..a317f54c35 100644 --- a/src/ophyd_async/epics/_backend/common.py +++ b/src/ophyd_async/epics/_backend/common.py @@ -1,6 +1,9 @@ +import inspect from enum import Enum from typing import Dict, Optional, Tuple, Type, TypedDict +from ophyd_async.core.signal_backend import RuntimeSubsetEnum + common_meta = { "units", "precision", @@ -30,19 +33,30 @@ def get_supported_values( datatype: Optional[Type[str]], pv_choices: Tuple[str, ...], ) -> Dict[str, str]: - if not datatype: + if inspect.isclass(datatype) and issubclass(datatype, RuntimeSubsetEnum): + if not set(datatype.choices).issubset(set(pv_choices)): + raise TypeError( + f"{pv} has choices {pv_choices}, " + f"which is not a superset of {str(datatype)}." + ) return {x: x or "_" for x in pv_choices} + elif inspect.isclass(datatype) and issubclass(datatype, Enum): + if not issubclass(datatype, str): + raise TypeError( + f"{pv} is type Enum but {datatype} does not inherit from String." + ) - if not issubclass(datatype, str): - raise TypeError(f"{pv} is type Enum but doesn't inherit from String") - if issubclass(datatype, Enum): choices = tuple(v.value for v in datatype) if set(choices) != set(pv_choices): raise TypeError( - ( - f"{pv} has choices {pv_choices}, " - f"which do not match {datatype}, which has {choices}" - ) + f"{pv} has choices {pv_choices}, " + f"which do not match {datatype}, which has {choices}." ) - return {x: datatype(x) for x in pv_choices} - return {x: x for x in pv_choices} + return {x: datatype(x) if x else "_" for x in pv_choices} + elif datatype is None: + return {x: x or "_" for x in pv_choices} + + raise TypeError( + f"{pv} has choices {pv_choices}. " + "Use an Enum or SubsetEnum to represent this." + ) diff --git a/tests/core/test_soft_signal_backend.py b/tests/core/test_soft_signal_backend.py index d21101160d..66d0bb07dd 100644 --- a/tests/core/test_soft_signal_backend.py +++ b/tests/core/test_soft_signal_backend.py @@ -31,7 +31,7 @@ def string_d(value): def enum_d(value): - return {"dtype": "string", "shape": [], "choices": ["Aaa", "Bbb", "Ccc"]} + return {"dtype": "string", "shape": [], "choices": ("Aaa", "Bbb", "Ccc")} def waveform_d(value): diff --git a/tests/core/test_subset_enum.py b/tests/core/test_subset_enum.py new file mode 100644 index 0000000000..4bcdcaf770 --- /dev/null +++ b/tests/core/test_subset_enum.py @@ -0,0 +1,90 @@ +import pytest +from epicscorelibs.ca import dbr +from p4p import Value as P4PValue +from p4p.nt import NTEnum + +from ophyd_async.core import SubsetEnum +from ophyd_async.epics._backend._aioca import make_converter as aioca_make_converter +from ophyd_async.epics._backend._p4p import make_converter as p4p_make_converter +from ophyd_async.epics.signal.signal import epics_signal_rw + + +async def test_runtime_enum_behaviour(): + rt_enum = SubsetEnum["A", "B"] + + with pytest.raises(RuntimeError) as exc: + rt_enum() + assert str(exc.value) == "SubsetEnum cannot be instantiated" + + assert issubclass(rt_enum, SubsetEnum) + + # Our metaclass doesn't cache already created runtime enums, + # so we can't do this + assert not issubclass(rt_enum, SubsetEnum["A", "B"]) + assert not issubclass(rt_enum, SubsetEnum["B", "A"]) + assert rt_enum is not SubsetEnum["A", "B"] + assert rt_enum is not SubsetEnum["B", "A"] + + assert str(rt_enum) == "SubsetEnum['A', 'B']" + assert str(SubsetEnum) == "SubsetEnum" + + with pytest.raises(TypeError) as exc: + SubsetEnum["A", "B", "A"] + assert str(exc.value) == "Duplicate elements in runtime enum choices." + + +async def test_ca_runtime_enum_converter(): + class EpicsValue: + def __init__(self): + self.name = "test" + self.ok = (True,) + self.errorcode = 0 + self.datatype = dbr.DBR_ENUM + self.element_count = 1 + self.severity = 0 + self.status = 0 + self.raw_stamp = (0,) + self.timestamp = 0 + self.datetime = 0 + self.enums = ["A", "B", "C"] # More than the runtime enum + + epics_value = EpicsValue() + rt_enum = SubsetEnum["A", "B"] + converter = aioca_make_converter( + rt_enum, values={"READ_PV": epics_value, "WRITE_PV": epics_value} + ) + assert converter.choices == {"A": "A", "B": "B", "C": "C"} + assert set(rt_enum.choices).issubset(set(converter.choices.keys())) + + +async def test_pva_runtime_enum_converter(): + enum_type = NTEnum.buildType() + epics_value = P4PValue( + enum_type, + { + "value.choices": ["A", "B", "C"], + }, + ) + rt_enum = SubsetEnum["A", "B"] + converter = p4p_make_converter( + rt_enum, values={"READ_PV": epics_value, "WRITE_PV": epics_value} + ) + assert {"A", "B"}.issubset(set(converter.choices)) + + +async def test_runtime_enum_signal(): + signal_rw_pva = epics_signal_rw(SubsetEnum["A1", "B1"], "ca://RW_PV", name="signal") + signal_rw_ca = epics_signal_rw(SubsetEnum["A2", "B2"], "ca://RW_PV", name="signal") + await signal_rw_pva.connect(mock=True) + await signal_rw_ca.connect(mock=True) + await signal_rw_pva.get_value() == "A1" + await signal_rw_ca.get_value() == "A2" + await signal_rw_pva.set("B1") + await signal_rw_ca.set("B2") + await signal_rw_pva.get_value() == "B1" + await signal_rw_ca.get_value() == "B2" + + # Will accept string values even if they're not in the runtime enum + # Though type checking should compain + await signal_rw_pva.set("C1") # type: ignore + await signal_rw_ca.set("C2") # type: ignore diff --git a/tests/epics/demo/test_demo.py b/tests/epics/demo/test_demo.py index 234afc3a77..d4931c6e4e 100644 --- a/tests/epics/demo/test_demo.py +++ b/tests/epics/demo/test_demo.py @@ -258,7 +258,7 @@ async def test_read_sensor(mock_sensor: demo.Sensor): ] == demo.EnergyMode.low desc = (await mock_sensor.describe_configuration())["mock_sensor-mode"] assert desc["dtype"] == "string" - assert desc["choices"] == ["Low Energy", "High Energy"] # type: ignore + assert desc["choices"] == ("Low Energy", "High Energy") # type: ignore set_mock_value(mock_sensor.mode, demo.EnergyMode.high) assert (await mock_sensor.read_configuration())["mock_sensor-mode"][ "value" diff --git a/tests/epics/test_signals.py b/tests/epics/test_signals.py index 764c2c3644..c89518e051 100644 --- a/tests/epics/test_signals.py +++ b/tests/epics/test_signals.py @@ -20,6 +20,7 @@ from bluesky.protocols import DataKey, Reading from ophyd_async.core import SignalBackend, T, load_from_yaml, save_to_yaml +from ophyd_async.core.signal_backend import SubsetEnum from ophyd_async.core.utils import NotConnected from ophyd_async.epics._backend.common import LimitPair, Limits from ophyd_async.epics.signal._epics_transport import EpicsTransport @@ -202,6 +203,8 @@ class MyEnum(str, Enum): c = "Ccc" +MySubsetEnum = SubsetEnum["Aaa", "Bbb", "Ccc"] + _metadata: Dict[str, Dict[str, Dict[str, Any]]] = { "ca": { "bool": {"units": ANY, "limits": ANY}, @@ -244,7 +247,10 @@ def get_dtype(suffix: str) -> str: d = {"dtype": dtype, "shape": [len(value)] if dtype == "array" else []} if get_internal_dtype(suffix) == "enum": - d["choices"] = [e.value for e in type(value)] + if issubclass(type(value), Enum): + d["choices"] = [e.value for e in type(value)] + else: + d["choices"] = list(value.choices) d.update(_metadata[protocol].get(get_internal_dtype(suffix), {})) @@ -497,11 +503,27 @@ class EnumNoString(Enum): ", which has ('Aaa', 'B', 'Ccc')" ), ), + ( + rt_enum := SubsetEnum["Aaa", "B", "Ccc"], + "enum", + ( + "has choices ('Aaa', 'Bbb', 'Ccc'), " + # SubsetEnum string output isn't deterministic + f"which is not a superset of {str(rt_enum)}." + ), + ), (int, "str", "has type str not int"), (str, "float", "has type float not str"), (str, "stra", "has type [str] not str"), (int, "uint8a", "has type [uint8] not int"), - (float, "enum", "is type Enum but doesn't inherit from String"), + ( + float, + "enum", + ( + "has choices ('Aaa', 'Bbb', 'Ccc'). " + "Use an Enum or SubsetEnum to represent this." + ), + ), (npt.NDArray[np.int32], "float64a", "has type [float64] not [int32]"), ], ) @@ -711,16 +733,14 @@ async def test_str_enum_returns_enum(ioc: IOC): assert val == "Bbb" -async def test_str_returns_enum(ioc: IOC): - await ioc.make_backend(str, "enum") +async def test_runtime_enum_returns_str(ioc: IOC): + await ioc.make_backend(MySubsetEnum, "enum") pv_name = f"{ioc.protocol}://{PV_PREFIX}:{ioc.protocol}:enum" + sig = epics_signal_rw(MySubsetEnum, pv_name) - sig = epics_signal_rw(str, pv_name) await sig.connect() val = await sig.get_value() - assert val == MyEnum.b assert val == "Bbb" - assert val is not MyEnum.b async def test_signal_returns_units_and_precision(ioc: IOC):