Add SupersetEnum and use it in ADBaseDataType (#864)

* Add SupersetEnum, rework Enums to not inherit from another and adjust framework to change

* Correct supported values check

* Organise imports and correct string docs

* Update src/ophyd_async/core/_signal_backend.py

grammar correction

Co-authored-by: Joseph Ware <[email protected]>

* Update src/ophyd_async/core/_signal_backend.py

grammar correction

Co-authored-by: Joseph Ware <[email protected]>

* Added tests and improved handling

* Fixed linting issue

* Improved test coverage and updated check for SupersetEnum

* Update src/ophyd_async/core/_signal_backend.py

Co-authored-by: Tom C (DLS) <[email protected]>

* Added UNDEFINED to ADBaseDataType

* Fix import

* Fix liniting issue and circular imports

* Further linting correction

* Added error message for UNDEFINED ADBaseDataType

* Potential fix for code scanning alert no. 186: First parameter of a class method is not named 'cls'

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

* Applied feedback

* Removed empty string values in datatypes

* Reverted suggest support_values check as broke tests

* Wrote SupersetEnum check in a more succinct way

Had to widen some types to do this

* Remove unused code

---------

Co-authored-by: Joseph Ware <[email protected]>
Co-authored-by: Tom C (DLS) <[email protected]>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: Tom Cobb <[email protected]>

Author: 4 people

src/ophyd_async/core/__init__.py

@@ -75,11 +75,13 @@
     DEFAULT_TIMEOUT,
     CalculatableTimeout,
     Callback,
+    EnumTypes,
     LazyMock,
     NotConnected,
     Reference,
     StrictEnum,
     SubsetEnum,
+    SupersetEnum,
     WatcherUpdate,
     gather_dict,
     get_dtype,
@@ -122,6 +124,8 @@
     "Array1D",
     "StrictEnum",
     "SubsetEnum",
+    "SupersetEnum",
+    "EnumTypes",
     "Table",
     "SignalMetadata",
     # Soft signal

src/ophyd_async/core/_signal_backend.py

@@ -6,8 +6,16 @@
 from bluesky.protocols import Reading
 from event_model import DataKey, Dtype, Limits
 
+from ophyd_async.core._utils import (
+    Callback,
+    EnumTypes,
+    StrictEnum,
+    SubsetEnum,
+    SupersetEnum,
+    get_enum_cls,
+)
+
 from ._table import Table
-from ._utils import Callback, StrictEnum, get_enum_cls
 
 DTypeScalar_co = TypeVar("DTypeScalar_co", covariant=True, bound=np.generic)
 """A numpy dtype like [](#numpy.float64)."""
@@ -24,7 +32,7 @@
 Primitive = bool | int | float | str
 SignalDatatype = (
     Primitive
-    | StrictEnum
+    | EnumTypes
     | Array1D[np.bool_]
     | Array1D[np.int8]
     | Array1D[np.uint8]
@@ -39,16 +47,18 @@
     | np.ndarray
     | Sequence[str]
     | Sequence[StrictEnum]
+    | Sequence[SubsetEnum]
+    | Sequence[SupersetEnum]
     | Table
 )
 """The supported [](#Signal) datatypes:
 
 - A python primitive [](#bool), [](#int), [](#float), [](#str)
-- A [](#StrictEnum) or [](#SubsetEnum) subclass
+- An [](#EnumTypes) subclass
 - A fixed datatype [](#Array1D) of numpy bool, signed and unsigned integers or float
 - A [](#numpy.ndarray) which can change dimensions and datatype at runtime
 - A sequence of [](#str)
-- A sequence of [](#StrictEnum) or [](#SubsetEnum) subclass
+- A sequence of [](#EnumTypes) subclasses
 - A [](#Table) subclass
 """
 # TODO: These typevars will not be needed when we drop python 3.11
@@ -58,7 +68,7 @@
 SignalDatatypeT = TypeVar("SignalDatatypeT", bound=SignalDatatype)
 """A typevar for a [](#SignalDatatype)."""
 SignalDatatypeV = TypeVar("SignalDatatypeV", bound=SignalDatatype)
-EnumT = TypeVar("EnumT", bound=StrictEnum)
+EnumT = TypeVar("EnumT", bound=EnumTypes)
 TableT = TypeVar("TableT", bound=Table)
 
 
@@ -136,7 +146,7 @@ def _datakey_dtype(datatype: type[SignalDatatype]) -> Dtype:
         or issubclass(datatype, Table)
     ):
         return "array"
-    elif issubclass(datatype, StrictEnum):
+    elif issubclass(datatype, EnumTypes):
         return "string"
     elif issubclass(datatype, Primitive):
         return _primitive_dtype[datatype]
@@ -153,7 +163,7 @@ def _datakey_dtype_numpy(
     elif (
         get_origin(datatype) is Sequence
         or datatype is str
-        or issubclass(datatype, StrictEnum)
+        or issubclass(datatype, EnumTypes)
     ):
         # TODO: use np.dtypes.StringDType when we can use in structured arrays
         # https://github.com/numpy/numpy/issues/25693
@@ -167,7 +177,7 @@ def _datakey_dtype_numpy(
 
 
 def _datakey_shape(value: SignalDatatype) -> list[int]:
-    if type(value) in _primitive_dtype or isinstance(value, StrictEnum):
+    if type(value) in _primitive_dtype or isinstance(value, EnumTypes):
         return []
     elif isinstance(value, np.ndarray):
         return list(value.shape)

src/ophyd_async/core/_utils.py

@@ -5,7 +5,15 @@
 from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass
 from enum import Enum, EnumMeta
-from typing import Any, Generic, Literal, ParamSpec, TypeVar, get_args, get_origin
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    ParamSpec,
+    TypeVar,
+    get_args,
+    get_origin,
+)
 from unittest.mock import Mock
 
 import numpy as np
@@ -19,20 +27,16 @@
 logger = logging.getLogger("ophyd_async")
 
 
-class StrictEnumMeta(EnumMeta):
-    def __new__(metacls, *args, **kwargs):
-        ret = super().__new__(metacls, *args, **kwargs)
+class UppercaseNameEnumMeta(EnumMeta):
+    def __new__(cls, *args, **kwargs):
+        ret = super().__new__(cls, *args, **kwargs)
         lowercase_names = [x.name for x in ret if not x.name.isupper()]  # type: ignore
         if lowercase_names:
             raise TypeError(f"Names {lowercase_names} should be uppercase")
         return ret
 
 
-class StrictEnum(str, Enum, metaclass=StrictEnumMeta):
-    """All members should exist in the Backend, and there will be no extras."""
-
-
-class SubsetEnumMeta(StrictEnumMeta):
+class AnyStringUppercaseNameEnumMeta(UppercaseNameEnumMeta):
     def __call__(self, value, *args, **kwargs):  # type: ignore
         """Return given value if it is a string and not a member of the enum.
 
@@ -54,10 +58,21 @@ def __call__(self, value, *args, **kwargs):  # type: ignore
         return super().__call__(value, *args, **kwargs)
 
 
-class SubsetEnum(StrictEnum, metaclass=SubsetEnumMeta):
+class StrictEnum(str, Enum, metaclass=UppercaseNameEnumMeta):
+    """All members should exist in the Backend, and there will be no extras."""
+
+
+class SubsetEnum(str, Enum, metaclass=AnyStringUppercaseNameEnumMeta):
     """All members should exist in the Backend, but there may be extras."""
 
 
+class SupersetEnum(str, Enum, metaclass=UppercaseNameEnumMeta):
+    """Some members should exist in the Backend, and there should be no extras."""
+
+
+EnumTypes = StrictEnum | SubsetEnum | SupersetEnum
+
+
 CALCULATE_TIMEOUT = "CALCULATE_TIMEOUT"
 """Sentinel used to implement ``myfunc(timeout=CalculateTimeout)``
 
@@ -207,10 +222,11 @@ def get_dtype(datatype: type) -> np.dtype:
     return np.dtype(get_args(get_args(datatype)[1])[0])
 
 
-def get_enum_cls(datatype: type | None) -> type[StrictEnum] | None:
+def get_enum_cls(datatype: type | None) -> type[EnumTypes] | None:
     """Get the enum class from a datatype.
 
-    :raises TypeError: if type is not a [](#StrictEnum) or [](#SubsetEnum) subclass
+    :raises TypeError: if type is not a [](#StrictEnum) or [](#SubsetEnum)
+    or [](#SupersetEnum) subclass
     ```python
     >>> from ophyd_async.core import StrictEnum
     >>> from collections.abc import Sequence
@@ -227,10 +243,11 @@ def get_enum_cls(datatype: type | None) -> type[StrictEnum] | None:
     if get_origin(datatype) is Sequence:
         datatype = get_args(datatype)[0]
     if datatype and issubclass(datatype, Enum):
-        if not issubclass(datatype, StrictEnum):
+        if not issubclass(datatype, EnumTypes):
             raise TypeError(
                 f"{datatype} should inherit from ophyd_async.core.SubsetEnum "
-                "or ophyd_async.core.StrictEnum"
+                "or ophyd_async.core.StrictEnum "
+                "or ophyd_async.core.SupersetEnum."
             )
         return datatype
     return None

src/ophyd_async/epics/adcore/_utils.py

@@ -7,11 +7,12 @@
     SignalRW,
     StrictEnum,
     SubsetEnum,
+    SupersetEnum,
     wait_for_value,
 )
 
 
-class ADBaseDataType(StrictEnum):
+class ADBaseDataType(SupersetEnum):
     INT8 = "Int8"
     UINT8 = "UInt8"
     INT16 = "Int16"
@@ -22,6 +23,9 @@ class ADBaseDataType(StrictEnum):
     UINT64 = "UInt64"
     FLOAT32 = "Float32"
     FLOAT64 = "Float64"
+    # Driver database override will blank the enum string if it doesn't
+    # support a datatype
+    UNDEFINED = ""
 
 
 def convert_ad_dtype_to_np(ad_dtype: ADBaseDataType) -> str:
@@ -37,7 +41,12 @@ def convert_ad_dtype_to_np(ad_dtype: ADBaseDataType) -> str:
         ADBaseDataType.FLOAT32: "<f4",
         ADBaseDataType.FLOAT64: "<f8",
     }
-    return ad_dtype_to_np_dtype[ad_dtype]
+    np_type = ad_dtype_to_np_dtype.get(ad_dtype)
+    if np_type is None:
+        raise ValueError(
+            "Areadetector driver has a blank DataType, this is not supported"
+        )
+    return np_type
 
 
 def convert_pv_dtype_to_np(datatype: str) -> str:

src/ophyd_async/epics/core/_aioca.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 import typing
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from functools import cache
 from math import isnan, nan
 from typing import Any, Generic, cast
@@ -146,7 +146,7 @@ def value(self, value: AugmentedValue) -> bool:
 
 
 class CaEnumConverter(CaConverter[str]):
-    def __init__(self, supported_values: dict[str, str]):
+    def __init__(self, supported_values: Mapping[str, str]):
         self.supported_values = supported_values
         super().__init__(
             str, dbr.DBR_STRING, metadata=SignalMetadata(choices=list(supported_values))

src/ophyd_async/epics/core/_util.py

@@ -1,16 +1,20 @@
-from collections.abc import Sequence
-from typing import Any, get_args, get_origin
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeVar, get_args, get_origin
 
 import numpy as np
 
 from ophyd_async.core import (
     SignalBackend,
     SignalDatatypeT,
+    StrictEnum,
     SubsetEnum,
+    SupersetEnum,
     get_dtype,
     get_enum_cls,
 )
 
+T = TypeVar("T")
+
 
 def get_pv_basename_and_field(pv: str) -> tuple[str, str | None]:
     """Split PV into record name and field."""
@@ -23,26 +27,30 @@ def get_pv_basename_and_field(pv: str) -> tuple[str, str | None]:
 
 def get_supported_values(
     pv: str,
-    datatype: type,
+    datatype: type[T],
     pv_choices: Sequence[str],
-) -> dict[str, str]:
+) -> Mapping[str, T | str]:
     enum_cls = get_enum_cls(datatype)
     if not enum_cls:
         raise TypeError(f"{datatype} is not an Enum")
     choices = [v.value for v in enum_cls]
     error_msg = f"{pv} has choices {pv_choices}, but {datatype} requested {choices} "
-    if issubclass(enum_cls, SubsetEnum):
+    if issubclass(enum_cls, StrictEnum):
+        if set(choices) != set(pv_choices):
+            raise TypeError(error_msg + "to be strictly equal to them.")
+    elif issubclass(enum_cls, SubsetEnum):
         if not set(choices).issubset(pv_choices):
             raise TypeError(error_msg + "to be a subset of them.")
+    elif issubclass(enum_cls, SupersetEnum):
+        if not set(pv_choices).issubset(choices):
+            raise TypeError(error_msg + "to be a superset of them.")
     else:
-        if set(choices) != set(pv_choices):
-            raise TypeError(error_msg + "to be strictly equal to them.")
-
-    # Take order from the pv choices
-    supported_values = {x: x for x in pv_choices}
-    # But override those that we specify via the datatype
-    for v in enum_cls:
-        supported_values[v.value] = v
+        raise TypeError(f"{datatype} is not a StrictEnum, SubsetEnum, or SupersetEnum")
+    # Create a map from the string value to the enum instance
+    # For StrictEnum and SupersetEnum, all values here will be enum values
+    # For SubsetEnum, only the values in choices will be enum values, the rest will be
+    # strings
+    supported_values = {x: enum_cls(x) for x in pv_choices}
     return supported_values
 
 

tests/epics/signal/test_common.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from ophyd_async.core import StrictEnum
+from ophyd_async.core import StrictEnum, SupersetEnum
 from ophyd_async.epics.core._util import get_supported_values  # noqa: PLC2701
 
 
@@ -45,3 +45,41 @@ class MyEnum(StrictEnum):
     assert len(supported_vals) == 2
     assert "test_1" in supported_vals
     assert "test_2" in supported_vals
+
+
+def test_given_supersetenum_that_partial_matches_are_valid():
+    class MyEnum(SupersetEnum):
+        TEST_1 = "test_1"
+        TEST_2 = "test_2"
+
+    supported_vals = get_supported_values("", MyEnum, ("test_1",))
+
+    assert "test_1" in supported_vals
+    assert supported_vals.get("test_2") is None
+
+
+class MySupersetEnum(SupersetEnum):
+    TEST_1 = "test_1"
+    TEST_2 = "test_2"
+
+
+SUPERSETENUM_ERROR_MESSAGE = "to be a superset of them."
+
+
+def test_given_supersetenum_that_all_values_plus_extra_values_are_invalid():
+    with pytest.raises(TypeError, match=SUPERSETENUM_ERROR_MESSAGE):
+        get_supported_values(
+            "",
+            MySupersetEnum,
+            (MySupersetEnum.TEST_1, MySupersetEnum.TEST_2, "extra_1"),
+        )
+
+
+def test_given_supersetenum_that_partial_values_plus_extra_values_are_invalid():
+    with pytest.raises(TypeError, match=SUPERSETENUM_ERROR_MESSAGE):
+        get_supported_values("", MySupersetEnum, (MySupersetEnum.TEST_1, "extra_1"))
+
+
+def test_given_supersetenum_that_no_matches_is_invalid():
+    with pytest.raises(TypeError, match=SUPERSETENUM_ERROR_MESSAGE):
+        get_supported_values("", MySupersetEnum, ("no_match_1", "no_match_2"))