ophyd-async 0.7.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

ophyd_async/core/_signal_backend.py

@@ -1,97 +1,164 @@
 from abc import abstractmethod
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    ClassVar,
-    Generic,
-    Literal,
-)
+from collections.abc import Sequence
+from typing import Generic, TypedDict, TypeVar, get_origin
 
+import numpy as np
 from bluesky.protocols import Reading
-from event_model import DataKey
-
-from ._utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
+from event_model import DataKey, Dtype, Limits
+
+from ._table import Table
+from ._utils import Callback, StrictEnum, T
+
+DTypeScalar_co = TypeVar("DTypeScalar_co", covariant=True, bound=np.generic)
+Array1D = np.ndarray[tuple[int], np.dtype[DTypeScalar_co]]
+Primitive = bool | int | float | str
+# NOTE: if you change this union then update the docs to match
+SignalDatatype = (
+    Primitive
+    | Array1D[np.bool_]
+    | Array1D[np.int8]
+    | Array1D[np.uint8]
+    | Array1D[np.int16]
+    | Array1D[np.uint16]
+    | Array1D[np.int32]
+    | Array1D[np.uint32]
+    | Array1D[np.int64]
+    | Array1D[np.uint64]
+    | Array1D[np.float32]
+    | Array1D[np.float64]
+    | np.ndarray
+    | StrictEnum
+    | Sequence[str]
+    | Sequence[StrictEnum]
+    | Table
+)
+# TODO: These typevars will not be needed when we drop python 3.11
+# as you can do MyConverter[SignalType: SignalTypeUnion]:
+# rather than MyConverter(Generic[SignalType])
+PrimitiveT = TypeVar("PrimitiveT", bound=Primitive)
+SignalDatatypeT = TypeVar("SignalDatatypeT", bound=SignalDatatype)
+SignalDatatypeV = TypeVar("SignalDatatypeV", bound=SignalDatatype)
+EnumT = TypeVar("EnumT", bound=StrictEnum)
+TableT = TypeVar("TableT", bound=Table)
 
 
-class SignalBackend(Generic[T]):
+class SignalBackend(Generic[SignalDatatypeT]):
     """A read/write/monitor backend for a Signals"""
 
-    #: Datatype of the signal value
-    datatype: type[T] | None = None
+    def __init__(self, datatype: type[SignalDatatypeT] | None):
+        self.datatype = datatype
 
-    @classmethod
     @abstractmethod
-    def datatype_allowed(cls, dtype: Any) -> bool:
-        """Check if a given datatype is acceptable for this signal backend."""
+    def source(self, name: str, read: bool) -> str:
+        """Return source of signal.
 
-    #: Like ca://PV_PREFIX:SIGNAL
-    @abstractmethod
-    def source(self, name: str) -> str:
-        """Return source of signal. Signals may pass a name to the backend, which can be
-        used or discarded."""
+        Signals may pass a name to the backend, which can be used or discarded.
+        """
 
     @abstractmethod
-    async def connect(self, timeout: float = DEFAULT_TIMEOUT):
+    async def connect(self, timeout: float):
         """Connect to underlying hardware"""
 
     @abstractmethod
-    async def put(self, value: T | None, wait=True, timeout=None):
-        """Put a value to the PV, if wait then wait for completion for up to timeout"""
+    async def put(self, value: SignalDatatypeT | None, wait: bool):
+        """Put a value to the PV, if wait then wait for completion"""
 
     @abstractmethod
     async def get_datakey(self, source: str) -> DataKey:
         """Metadata like source, dtype, shape, precision, units"""
 
     @abstractmethod
-    async def get_reading(self) -> Reading:
+    async def get_reading(self) -> Reading[SignalDatatypeT]:
         """The current value, timestamp and severity"""
 
     @abstractmethod
-    async def get_value(self) -> T:
+    async def get_value(self) -> SignalDatatypeT:
         """The current value"""
 
     @abstractmethod
-    async def get_setpoint(self) -> T:
+    async def get_setpoint(self) -> SignalDatatypeT:
         """The point that a signal was requested to move to."""
 
     @abstractmethod
-    def set_callback(self, callback: ReadingValueCallback[T] | None) -> None:
+    def set_callback(self, callback: Callback[T] | None) -> 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)}"  # type: ignore
-        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
+_primitive_dtype: dict[type[Primitive], Dtype] = {
+    bool: "boolean",
+    int: "integer",
+    float: "number",
+    str: "string",
+}
+
+
+class SignalMetadata(TypedDict, total=False):
+    limits: Limits
+    choices: list[str]
+    precision: int
+    units: str
+
+
+def _datakey_dtype(datatype: type[SignalDatatype]) -> Dtype:
+    if (
+        datatype is np.ndarray
+        or get_origin(datatype) in (Sequence, np.ndarray)
+        or issubclass(datatype, Table)
+    ):
+        return "array"
+    elif issubclass(datatype, StrictEnum):
+        return "string"
+    elif issubclass(datatype, Primitive):
+        return _primitive_dtype[datatype]
+    else:
+        raise TypeError(f"Can't make dtype for {datatype}")
+
+
+def _datakey_dtype_numpy(
+    datatype: type[SignalDatatypeT], value: SignalDatatypeT
+) -> np.dtype:
+    if isinstance(value, np.ndarray):
+        # The value already has a dtype, use that
+        return value.dtype
+    elif (
+        get_origin(datatype) is Sequence
+        or datatype is str
+        or issubclass(datatype, StrictEnum)
+    ):
+        # TODO: use np.dtypes.StringDType when we can use in structured arrays
+        # https://github.com/numpy/numpy/issues/25693
+        return np.dtype("S40")
+    elif isinstance(value, Table):
+        return value.numpy_dtype()
+    elif issubclass(datatype, Primitive):
+        return np.dtype(datatype)
+    else:
+        raise TypeError(f"Can't make dtype_numpy for {datatype}")
+
+
+def _datakey_shape(value: SignalDatatype) -> list[int]:
+    if type(value) in _primitive_dtype or isinstance(value, StrictEnum):
+        return []
+    elif isinstance(value, np.ndarray):
+        return list(value.shape)
+    elif isinstance(value, Sequence | Table):
+        return [len(value)]
+    else:
+        raise TypeError(f"Can't make shape for {value}")
+
+
+def make_datakey(
+    datatype: type[SignalDatatypeT],
+    value: SignalDatatypeT,
+    source: str,
+    metadata: SignalMetadata,
+) -> DataKey:
+    dtn = _datakey_dtype_numpy(datatype, value)
+    return DataKey(
+        dtype=_datakey_dtype(datatype),
+        shape=_datakey_shape(value),
+        # Ignore until https://github.com/bluesky/event-model/issues/308
+        dtype_numpy=dtn.descr if len(dtn.descr) > 1 else dtn.str,  # type: ignore
+        source=source,
+        **metadata,
+    )

ophyd_async/core/_soft_signal_backend.py

@@ -1,244 +1,181 @@
 from __future__ import annotations
 
-import inspect
 import time
-from collections import abc
-from enum import Enum
-from typing import Generic, cast, get_origin
+from abc import abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Generic, get_args, get_origin
 
 import numpy as np
 from bluesky.protocols import Reading
 from event_model import DataKey
-from event_model.documents.event_descriptor import Dtype
-from pydantic import BaseModel
-from typing_extensions import TypedDict
 
 from ._signal_backend import (
-    RuntimeSubsetEnum,
+    Array1D,
+    EnumT,
+    Primitive,
+    PrimitiveT,
     SignalBackend,
+    SignalDatatype,
+    SignalDatatypeT,
+    SignalMetadata,
+    TableT,
+    make_datakey,
 )
-from ._utils import (
-    DEFAULT_TIMEOUT,
-    ReadingValueCallback,
-    T,
-    get_dtype,
-    is_pydantic_model,
-)
-
-primitive_dtypes: dict[type, Dtype] = {
-    str: "string",
-    int: "integer",
-    float: "number",
-    bool: "boolean",
-}
+from ._table import Table
+from ._utils import Callback, get_dtype, get_enum_cls
 
 
-class SignalMetadata(TypedDict):
-    units: str | None
-    precision: int | None
+class SoftConverter(Generic[SignalDatatypeT]):
+    # This is Any -> SignalDatatypeT because we support coercing
+    # value types to SignalDatatype to allow people to do things like
+    # SignalRW[Enum].set("enum value")
+    @abstractmethod
+    def write_value(self, value: Any) -> SignalDatatypeT: ...
 
 
-class SoftConverter(Generic[T]):
-    def value(self, value: T) -> T:
-        return value
+@dataclass
+class PrimitiveSoftConverter(SoftConverter[PrimitiveT]):
+    datatype: type[PrimitiveT]
 
-    def write_value(self, value: T) -> T:
-        return value
+    def write_value(self, value: Any) -> PrimitiveT:
+        return self.datatype(value) if value else self.datatype()
 
-    def reading(self, value: T, timestamp: float, severity: int) -> Reading:
-        return Reading(
-            value=value,
-            timestamp=timestamp,
-            alarm_severity=-1 if severity > 2 else severity,
-        )
 
-    def get_datakey(self, source: str, value, **metadata) -> DataKey:
-        dk: DataKey = {"source": source, "shape": [], **metadata}  # type: ignore
-        dtype = type(value)
-        if np.issubdtype(dtype, np.integer):
-            dtype = int
-        elif np.issubdtype(dtype, np.floating):
-            dtype = float
-        assert (
-            dtype in primitive_dtypes
-        ), f"invalid converter for value of type {type(value)}"
-        dk["dtype"] = primitive_dtypes[dtype]
-        # type ignore until https://github.com/bluesky/event-model/issues/308
-        try:
-            dk["dtype_numpy"] = np.dtype(dtype).descr[0][1]  # type: ignore
-        except TypeError:
-            dk["dtype_numpy"] = ""  # type: ignore
-        return dk
-
-    def make_initial_value(self, datatype: type[T] | None) -> T:
-        if datatype is None:
-            return cast(T, None)
-
-        return datatype()
-
-
-class SoftArrayConverter(SoftConverter):
-    def get_datakey(self, source: str, value, **metadata) -> DataKey:
-        dtype_numpy = ""
-        if isinstance(value, list):
-            if len(value) > 0:
-                dtype_numpy = np.dtype(type(value[0])).descr[0][1]
-        else:
-            dtype_numpy = np.dtype(value.dtype).descr[0][1]
+class SequenceStrSoftConverter(SoftConverter[Sequence[str]]):
+    def write_value(self, value: Any) -> Sequence[str]:
+        return [str(v) for v in value] if value else []
 
-        return {
-            "source": source,
-            "dtype": "array",
-            "dtype_numpy": dtype_numpy,  # type: ignore
-            "shape": [len(value)],
-            **metadata,
-        }
 
-    def make_initial_value(self, datatype: type[T] | None) -> T:
-        if datatype is None:
-            return cast(T, None)
+@dataclass
+class SequenceEnumSoftConverter(SoftConverter[Sequence[EnumT]]):
+    datatype: type[EnumT]
 
-        if get_origin(datatype) == abc.Sequence:
-            return cast(T, [])
+    def write_value(self, value: Any) -> Sequence[EnumT]:
+        return [self.datatype(v) for v in value] if value else []
 
-        return cast(T, datatype(shape=0))  # type: ignore
 
+@dataclass
+class NDArraySoftConverter(SoftConverter[Array1D]):
+    datatype: np.dtype | None = None
 
-class SoftEnumConverter(SoftConverter):
-    choices: tuple[str, ...]
-
-    def __init__(self, datatype: RuntimeSubsetEnum | type[Enum]):
-        if issubclass(datatype, Enum):  # type: ignore
-            self.choices = tuple(v.value for v in datatype)
-        else:
-            self.choices = datatype.choices
-
-    def write_value(self, value: Enum | str) -> str:
-        return value  # type: ignore
+    def write_value(self, value: Any) -> Array1D:
+        return np.array(() if value is None else value, dtype=self.datatype)
 
-    def get_datakey(self, source: str, value, **metadata) -> DataKey:
-        return {
-            "source": source,
-            "dtype": "string",
-            # type ignore until https://github.com/bluesky/event-model/issues/308
-            "dtype_numpy": "|S40",  # type: ignore
-            "shape": [],
-            "choices": self.choices,
-            **metadata,
-        }
 
-    def make_initial_value(self, datatype: type[T] | None) -> T:
-        if datatype is None:
-            return cast(T, None)
+@dataclass
+class EnumSoftConverter(SoftConverter[EnumT]):
+    datatype: type[EnumT]
 
-        if issubclass(datatype, Enum):
-            return cast(T, list(datatype.__members__.values())[0])  # type: ignore
-        return cast(T, self.choices[0])
+    def write_value(self, value: Any) -> EnumT:
+        return (
+            self.datatype(value)
+            if value
+            else list(self.datatype.__members__.values())[0]
+        )
 
 
-class SoftPydanticModelConverter(SoftConverter):
-    def __init__(self, datatype: type[BaseModel]):
-        self.datatype = datatype
+@dataclass
+class TableSoftConverter(SoftConverter[TableT]):
+    datatype: type[TableT]
 
-    def write_value(self, value):
+    def write_value(self, value: Any) -> TableT:
         if isinstance(value, dict):
             return self.datatype(**value)
-        return value
-
-
-def make_converter(datatype):
-    is_array = get_dtype(datatype) is not None
-    is_sequence = get_origin(datatype) == abc.Sequence
-    is_enum = inspect.isclass(datatype) and (
-        issubclass(datatype, Enum) or issubclass(datatype, RuntimeSubsetEnum)
-    )
-
-    if is_array or is_sequence:
-        return SoftArrayConverter()
-    if is_enum:
-        return SoftEnumConverter(datatype)  # type: ignore
-    if is_pydantic_model(datatype):
-        return SoftPydanticModelConverter(datatype)  # type: ignore
-
-    return SoftConverter()
-
-
-class SoftSignalBackend(SignalBackend[T]):
+        elif isinstance(value, self.datatype):
+            return value
+        elif value is None:
+            return self.datatype()
+        else:
+            raise TypeError(f"Cannot convert {value} to {self.datatype}")
+
+
+@lru_cache
+def make_converter(datatype: type[SignalDatatype]) -> SoftConverter:
+    enum_cls = get_enum_cls(datatype)
+    if datatype == Sequence[str]:
+        return SequenceStrSoftConverter()
+    elif get_origin(datatype) == Sequence and enum_cls:
+        return SequenceEnumSoftConverter(enum_cls)
+    elif datatype is np.ndarray:
+        return NDArraySoftConverter()
+    elif get_origin(datatype) == np.ndarray:
+        if datatype not in get_args(SignalDatatype):
+            raise TypeError(f"Expected Array1D[dtype], got {datatype}")
+        return NDArraySoftConverter(get_dtype(datatype))
+    elif enum_cls:
+        return EnumSoftConverter(enum_cls)
+    elif issubclass(datatype, Table):
+        return TableSoftConverter(datatype)
+    elif issubclass(datatype, Primitive):
+        return PrimitiveSoftConverter(datatype)
+    raise TypeError(f"Can't make converter for {datatype}")
+
+
+class SoftSignalBackend(SignalBackend[SignalDatatypeT]):
     """An backend to a soft Signal, for test signals see ``MockSignalBackend``."""
 
-    _value: T
-    _initial_value: T | None
-    _timestamp: float
-    _severity: int
-
-    @classmethod
-    def datatype_allowed(cls, dtype: type) -> bool:
-        return True  # Any value allowed in a soft signal
-
     def __init__(
         self,
-        datatype: type[T] | None,
-        initial_value: T | None = None,
-        metadata: SignalMetadata = None,  # type: ignore
-    ) -> None:
-        self.datatype = datatype
-        self._initial_value = initial_value
-        self._metadata = metadata or {}
-        self.converter: SoftConverter = make_converter(datatype)
-        if self._initial_value is None:
-            self._initial_value = self.converter.make_initial_value(self.datatype)
-        else:
-            self._initial_value = self.converter.write_value(self._initial_value)  # type: ignore
-
-        self.callback: ReadingValueCallback[T] | None = None
-        self._severity = 0
-        self.set_value(self._initial_value)  # type: ignore
+        datatype: type[SignalDatatypeT] | None,
+        initial_value: SignalDatatypeT | None = None,
+        units: str | None = None,
+        precision: int | None = None,
+    ):
+        # Create the right converter for the datatype
+        self.converter = make_converter(datatype or float)
+        # Add the extra static metadata to the dictionary
+        self.metadata: SignalMetadata = {}
+        if units is not None:
+            self.metadata["units"] = units
+        if precision is not None:
+            self.metadata["precision"] = precision
+        if enum_cls := get_enum_cls(datatype):
+            self.metadata["choices"] = [v.value for v in enum_cls]
+        # Create and set the initial value
+        self.initial_value = self.converter.write_value(initial_value)
+        self.reading: Reading[SignalDatatypeT]
+        self.callback: Callback[Reading[SignalDatatypeT]] | None = None
+        self.set_value(self.initial_value)
+        super().__init__(datatype)
+
+    def set_value(self, value: SignalDatatypeT):
+        self.reading = Reading(
+            value=self.converter.write_value(value),
+            timestamp=time.monotonic(),
+            alarm_severity=0,
+        )
+        if self.callback:
+            self.callback(self.reading)
 
-    def source(self, name: str) -> str:
+    def source(self, name: str, read: bool) -> str:
         return f"soft://{name}"
 
-    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
-        """Connection isn't required for soft signals."""
+    async def connect(self, timeout: float):
         pass
 
-    async def put(self, value: T | None, wait=True, timeout=None):
-        write_value = (
-            self.converter.write_value(value)
-            if value is not None
-            else self._initial_value
-        )
-
-        self.set_value(write_value)  # type: ignore
-
-    def set_value(self, value: T):
-        """Method to bypass asynchronous logic."""
-        self._value = value
-        self._timestamp = time.monotonic()
-        reading: Reading = self.converter.reading(
-            self._value, self._timestamp, self._severity
-        )
-
-        if self.callback:
-            self.callback(reading, self._value)
+    async def put(self, value: SignalDatatypeT | None, wait: bool) -> None:
+        write_value = self.initial_value if value is None else value
+        self.set_value(write_value)
 
     async def get_datakey(self, source: str) -> DataKey:
-        return self.converter.get_datakey(source, self._value, **self._metadata)
+        return make_datakey(
+            self.datatype or float, self.reading["value"], source, self.metadata
+        )
 
-    async def get_reading(self) -> Reading:
-        return self.converter.reading(self._value, self._timestamp, self._severity)
+    async def get_reading(self) -> Reading[SignalDatatypeT]:
+        return self.reading
 
-    async def get_value(self) -> T:
-        return self.converter.value(self._value)
+    async def get_value(self) -> SignalDatatypeT:
+        return self.reading["value"]
 
-    async def get_setpoint(self) -> T:
-        """For a soft signal, the setpoint and readback values are the same."""
-        return await self.get_value()
+    async def get_setpoint(self) -> SignalDatatypeT:
+        # For a soft signal, the setpoint and readback values are the same.
+        return self.reading["value"]
 
-    def set_callback(self, callback: ReadingValueCallback[T] | None) -> None:
+    def set_callback(self, callback: Callback[Reading[SignalDatatypeT]] | None) -> None:
         if callback:
             assert not self.callback, "Cannot set a callback when one is already set"
-            reading: Reading = self.converter.reading(
-                self._value, self._timestamp, self._severity
-            )
-            callback(reading, self._value)
+            callback(self.reading)
         self.callback = callback