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

ophyd_async/epics/signal/_aioca.py

@@ -1,11 +1,8 @@
-import inspect
 import logging
 import sys
 from collections.abc import Sequence
-from dataclasses import dataclass
-from enum import Enum
 from math import isnan, nan
-from typing import Any, get_origin
+from typing import Any, Generic, cast
 
 import numpy as np
 from aioca import (
@@ -21,233 +18,201 @@ from aioca import (
 from aioca.types import AugmentedValue, Dbr, Format
 from bluesky.protocols import Reading
 from epicscorelibs.ca import dbr
-from event_model import DataKey
-from event_model.documents.event_descriptor import Dtype
+from event_model import DataKey, Limits, LimitsRange
 
 from ophyd_async.core import (
-    DEFAULT_TIMEOUT,
+    Array1D,
+    Callback,
     NotConnected,
-    ReadingValueCallback,
-    RuntimeSubsetEnum,
     SignalBackend,
-    T,
-    get_dtype,
+    SignalDatatype,
+    SignalDatatypeT,
+    SignalMetadata,
+    get_enum_cls,
     get_unique,
+    make_datakey,
     wait_for_connection,
 )
 
-from ._common import LimitPair, Limits, common_meta, get_supported_values
-
-dbr_to_dtype: dict[Dbr, Dtype] = {
-    dbr.DBR_STRING: "string",
-    dbr.DBR_SHORT: "integer",
-    dbr.DBR_FLOAT: "number",
-    dbr.DBR_CHAR: "string",
-    dbr.DBR_LONG: "integer",
-    dbr.DBR_DOUBLE: "number",
-}
-
-
-def _data_key_from_augmented_value(
-    value: AugmentedValue,
-    *,
-    choices: list[str] | None = None,
-    dtype: Dtype | None = None,
-) -> DataKey:
-    """Use the return value of get with FORMAT_CTRL to construct a DataKey
-    describing the signal. See docstring of AugmentedValue for expected
-    value fields by DBR type.
-
-    Args:
-        value (AugmentedValue): Description of the the return type of a DB record
-        choices: Optional list of enum choices to pass as metadata in the datakey
-        dtype: Optional override dtype when AugmentedValue is ambiguous, e.g. booleans
-
-    Returns:
-        DataKey: A rich DataKey describing the DB record
-    """
-    source = f"ca://{value.name}"
-    assert value.ok, f"Error reading {source}: {value}"
-
-    scalar = value.element_count == 1
-    dtype = dtype or dbr_to_dtype[value.datatype]  # type: ignore
-
-    dtype_numpy = np.dtype(dbr.DbrCodeToType[value.datatype].dtype).descr[0][1]
-
-    d = DataKey(
-        source=source,
-        dtype=dtype if scalar else "array",
-        # Ignore until https://github.com/bluesky/event-model/issues/308
-        dtype_numpy=dtype_numpy,  # type: ignore
-        # strictly value.element_count >= len(value)
-        shape=[] if scalar else [len(value)],
-    )
-    for key in common_meta:
-        attr = getattr(value, key, nan)
-        if isinstance(attr, str) or not isnan(attr):
-            d[key] = attr
-
-    if choices is not None:
-        d["choices"] = choices  # type: ignore
-
-    if limits := _limits_from_augmented_value(value):
-        d["limits"] = limits  # type: ignore
-
-    return d
+from ._common import format_datatype, get_supported_values
 
 
 def _limits_from_augmented_value(value: AugmentedValue) -> Limits:
-    def get_limits(limit: str) -> LimitPair:
+    def get_limits(limit: str) -> LimitsRange | None:
         low = getattr(value, f"lower_{limit}_limit", nan)
         high = getattr(value, f"upper_{limit}_limit", nan)
-        return LimitPair(
-            low=None if isnan(low) else low, high=None if isnan(high) else high
-        )
-
-    return Limits(
-        alarm=get_limits("alarm"),
-        control=get_limits("ctrl"),
-        display=get_limits("disp"),
-        warning=get_limits("warning"),
-    )
-
+        if not (isnan(low) and isnan(high)):
+            return LimitsRange(
+                low=None if isnan(low) else low,
+                high=None if isnan(high) else high,
+            )
 
-@dataclass
-class CaConverter:
-    read_dbr: Dbr | None
-    write_dbr: Dbr | None
+    limits = Limits()
+    if limits_range := get_limits("alarm"):
+        limits["alarm"] = limits_range
+    if limits_range := get_limits("ctrl"):
+        limits["control"] = limits_range
+    if limits_range := get_limits("disp"):
+        limits["display"] = limits_range
+    if limits_range := get_limits("warning"):
+        limits["warning"] = limits_range
+    return limits
+
+
+def _metadata_from_augmented_value(
+    value: AugmentedValue, metadata: SignalMetadata
+) -> SignalMetadata:
+    metadata = metadata.copy()
+    if hasattr(value, "units"):
+        metadata["units"] = value.units
+    if hasattr(value, "precision") and not isnan(value.precision):
+        metadata["precision"] = value.precision
+    if limits := _limits_from_augmented_value(value):
+        metadata["limits"] = limits
+    return metadata
+
+
+class CaConverter(Generic[SignalDatatypeT]):
+    def __init__(
+        self,
+        datatype: type[SignalDatatypeT],
+        read_dbr: Dbr,
+        write_dbr: Dbr | None = None,
+        metadata: SignalMetadata | None = None,
+    ):
+        self.datatype = datatype
+        self.read_dbr: Dbr = read_dbr
+        self.write_dbr: Dbr | None = write_dbr
+        self.metadata = metadata or SignalMetadata()
 
-    def write_value(self, value) -> Any:
+    def write_value(self, value: Any) -> Any:
+        # The ca library will do the conversion for us
         return value
 
-    def value(self, value: AugmentedValue):
+    def value(self, value: AugmentedValue) -> SignalDatatypeT:
         # for channel access ca_xxx classes, this
         # invokes __pos__ operator to return an instance of
         # the builtin base class
         return +value  # type: ignore
 
-    def reading(self, value: AugmentedValue) -> Reading:
-        return {
-            "value": self.value(value),
-            "timestamp": value.timestamp,
-            "alarm_severity": -1 if value.severity > 2 else value.severity,
-        }
-
-    def get_datakey(self, value: AugmentedValue) -> DataKey:
-        return _data_key_from_augmented_value(value)
 
-
-class CaLongStrConverter(CaConverter):
-    def __init__(self):
-        return super().__init__(dbr.DBR_CHAR_STR, dbr.DBR_CHAR_STR)
-
-    def write_value(self, value: str):
-        # Add a null in here as this is what the commandline caput does
-        # TODO: this should be in the server so check if it can be pushed to asyn
-        return value + "\0"
+class DisconnectedCaConverter(CaConverter):
+    def __getattribute__(self, __name: str) -> Any:
+        raise NotImplementedError("No PV has been set as connect() has not been called")
 
 
-class CaArrayConverter(CaConverter):
-    def value(self, value: AugmentedValue):
+class CaArrayConverter(CaConverter[np.ndarray]):
+    def value(self, value: AugmentedValue) -> np.ndarray:
+        # A less expensive conversion
         return np.array(value, copy=False)
 
 
-@dataclass
-class CaEnumConverter(CaConverter):
-    """To prevent issues when a signal is restarted and returns with different enum
-    values or orders, we put treat an Enum signal as a string, and cache the
-    choices on this class.
-    """
+class CaSequenceStrConverter(CaConverter[Sequence[str]]):
+    def value(self, value: AugmentedValue) -> Sequence[str]:
+        return [str(v) for v in value]  # type: ignore
 
-    choices: dict[str, str]
 
-    def write_value(self, value: Enum | str):
-        if isinstance(value, Enum):
-            return value.value
-        else:
-            return value
+class CaLongStrConverter(CaConverter[str]):
+    def __init__(self):
+        super().__init__(str, dbr.DBR_CHAR_STR, dbr.DBR_CHAR_STR)
 
-    def value(self, value: AugmentedValue):
-        return self.choices[value]  # type: ignore
+    def write_value_and_dbr(self, value: Any) -> Any:
+        # Add a null in here as this is what the commandline caput does
+        # TODO: this should be in the server so check if it can be pushed to asyn
+        return value + "\0"
 
-    def get_datakey(self, value: AugmentedValue) -> DataKey:
-        # Sometimes DBR_TYPE returns as String, must pass choices still
-        return _data_key_from_augmented_value(value, choices=list(self.choices.keys()))
 
+class CaBoolConverter(CaConverter[bool]):
+    def __init__(self):
+        super().__init__(bool, dbr.DBR_SHORT)
 
-@dataclass
-class CaBoolConverter(CaConverter):
     def value(self, value: AugmentedValue) -> bool:
         return bool(value)
 
-    def get_datakey(self, value: AugmentedValue) -> DataKey:
-        return _data_key_from_augmented_value(value, dtype="boolean")
 
+class CaEnumConverter(CaConverter[str]):
+    def __init__(self, supported_values: dict[str, str]):
+        self.supported_values = supported_values
+        super().__init__(
+            str, dbr.DBR_STRING, metadata=SignalMetadata(choices=list(supported_values))
+        )
 
-class DisconnectedCaConverter(CaConverter):
-    def __getattribute__(self, __name: str) -> Any:
-        raise NotImplementedError("No PV has been set as connect() has not been called")
+    def value(self, value: AugmentedValue) -> str:
+        return self.supported_values[str(value)]
+
+
+_datatype_converter_from_dbr: dict[
+    tuple[Dbr, bool], tuple[type[SignalDatatype], type[CaConverter]]
+] = {
+    (dbr.DBR_STRING, False): (str, CaConverter),
+    (dbr.DBR_SHORT, False): (int, CaConverter),
+    (dbr.DBR_FLOAT, False): (float, CaConverter),
+    (dbr.DBR_ENUM, False): (str, CaConverter),
+    (dbr.DBR_CHAR, False): (int, CaConverter),
+    (dbr.DBR_LONG, False): (int, CaConverter),
+    (dbr.DBR_DOUBLE, False): (float, CaConverter),
+    (dbr.DBR_STRING, True): (Sequence[str], CaSequenceStrConverter),
+    (dbr.DBR_SHORT, True): (Array1D[np.int16], CaArrayConverter),
+    (dbr.DBR_FLOAT, True): (Array1D[np.float32], CaArrayConverter),
+    (dbr.DBR_ENUM, True): (Sequence[str], CaSequenceStrConverter),
+    (dbr.DBR_CHAR, True): (Array1D[np.uint8], CaArrayConverter),
+    (dbr.DBR_LONG, True): (Array1D[np.int32], CaArrayConverter),
+    (dbr.DBR_DOUBLE, True): (Array1D[np.float64], CaArrayConverter),
+}
 
 
 def make_converter(
     datatype: type | None, values: dict[str, AugmentedValue]
 ) -> CaConverter:
     pv = list(values)[0]
-    pv_dbr = get_unique({k: v.datatype for k, v in values.items()}, "datatypes")
+    pv_dbr = cast(
+        Dbr, get_unique({k: v.datatype for k, v in values.items()}, "datatypes")
+    )
     is_array = bool([v for v in values.values() if v.element_count > 1])
-    if is_array and datatype is str and pv_dbr == dbr.DBR_CHAR:
+    # Infer a datatype and converter from the dbr
+    inferred_datatype, converter_cls = _datatype_converter_from_dbr[(pv_dbr, is_array)]
+    # Some override cases
+    if is_array and pv_dbr == dbr.DBR_CHAR and datatype is str:
         # Override waveform of chars to be treated as string
         return CaLongStrConverter()
-    elif is_array and pv_dbr == dbr.DBR_STRING:
-        # Waveform of strings, check we wanted this
-        if datatype:
-            datatype_dtype = get_dtype(datatype)
-            if not datatype_dtype or not np.can_cast(datatype_dtype, np.str_):
-                raise TypeError(f"{pv} has type [str] not {datatype.__name__}")
-        return CaArrayConverter(pv_dbr, None)
-    elif is_array:
-        pv_dtype = get_unique({k: v.dtype for k, v in values.items()}, "dtypes")  # type: ignore
-        # This is an array
-        if datatype:
-            # Check we wanted an array of this type
-            dtype = get_dtype(datatype)
-            if not dtype:
-                raise TypeError(f"{pv} has type [{pv_dtype}] not {datatype.__name__}")
-            if dtype != pv_dtype:
-                raise TypeError(f"{pv} has type [{pv_dtype}] not [{dtype}]")
-        return CaArrayConverter(pv_dbr, None)  # type: ignore
-    elif pv_dbr == dbr.DBR_ENUM and datatype is bool:
-        # Database can't do bools, so are often representated as enums,
-        # CA can do int
-        pv_choices_len = get_unique(
+    elif not is_array and datatype is bool and pv_dbr == dbr.DBR_ENUM:
+        # Database can't do bools, so are often representated as enums of len 2
+        pv_num_choices = get_unique(
             {k: len(v.enums) for k, v in values.items()}, "number of choices"
         )
-        if pv_choices_len != 2:
-            raise TypeError(f"{pv} has {pv_choices_len} choices, can't map to bool")
-        return CaBoolConverter(dbr.DBR_SHORT, dbr.DBR_SHORT)
-    elif pv_dbr == dbr.DBR_ENUM:
-        # This is an Enum
+        if pv_num_choices != 2:
+            raise TypeError(f"{pv} has {pv_num_choices} choices, can't map to bool")
+        return CaBoolConverter()
+    elif not is_array and pv_dbr == dbr.DBR_ENUM:
         pv_choices = get_unique(
             {k: tuple(v.enums) for k, v in values.items()}, "choices"
         )
-        supported_values = get_supported_values(pv, datatype, pv_choices)
-        return CaEnumConverter(dbr.DBR_STRING, None, supported_values)
-    else:
-        value = list(values.values())[0]
-        # Done the dbr check, so enough to check one of the values
-        if datatype and not isinstance(value, datatype):
-            # Allow int signals to represent float records when prec is 0
-            is_prec_zero_float = (
-                isinstance(value, float)
-                and get_unique({k: v.precision for k, v in values.items()}, "precision")
-                == 0
+        if enum_cls := get_enum_cls(datatype):
+            # If explicitly requested then check
+            return CaEnumConverter(get_supported_values(pv, enum_cls, pv_choices))
+        elif datatype in (None, str):
+            # Drop to string for safety, but retain choices as metadata
+            return CaConverter(
+                str,
+                dbr.DBR_STRING,
+                metadata=SignalMetadata(choices=list(pv_choices)),
             )
-            if not (datatype is int and is_prec_zero_float):
-                raise TypeError(
-                    f"{pv} has type {type(value).__name__.replace('ca_', '')} "
-                    + f"not {datatype.__name__}"
-                )
-    return CaConverter(pv_dbr, None)  # type: ignore
+    elif (
+        inferred_datatype is float
+        and datatype is int
+        and get_unique({k: v.precision for k, v in values.items()}, "precision") == 0
+    ):
+        # Allow int signals to represent float records when prec is 0
+        return CaConverter(int, pv_dbr)
+    elif datatype in (None, inferred_datatype):
+        # If datatype matches what we are given then allow it and use inferred converter
+        return converter_cls(inferred_datatype, pv_dbr)
+    if pv_dbr == dbr.DBR_ENUM:
+        inferred_datatype = "str | SubsetEnum | StrictEnum"
+    raise TypeError(
+        f"{pv} with inferred datatype {format_datatype(inferred_datatype)}"
+        f" cannot be coerced to {format_datatype(datatype)}"
+    )
 
 
 _tried_pyepics = False
@@ -262,42 +227,24 @@ def _use_pyepics_context_if_imported():
         _tried_pyepics = True
 
 
-class CaSignalBackend(SignalBackend[T]):
-    _ALLOWED_DATATYPES = (
-        bool,
-        int,
-        float,
-        str,
-        Sequence,
-        Enum,
-        RuntimeSubsetEnum,
-        np.ndarray,
-    )
-
-    @classmethod
-    def datatype_allowed(cls, dtype: Any) -> bool:
-        stripped_origin = get_origin(dtype) or dtype
-        if dtype is None:
-            return True
-
-        return inspect.isclass(stripped_origin) and issubclass(
-            stripped_origin, cls._ALLOWED_DATATYPES
-        )
-
-    def __init__(self, datatype: type[T] | None, read_pv: str, write_pv: str):
-        self.datatype = datatype
-        if not CaSignalBackend.datatype_allowed(self.datatype):
-            raise TypeError(f"Given datatype {self.datatype} unsupported in CA.")
+class CaSignalBackend(SignalBackend[SignalDatatypeT]):
+    def __init__(
+        self,
+        datatype: type[SignalDatatypeT] | None,
+        read_pv: str = "",
+        write_pv: str = "",
+    ):
         self.read_pv = read_pv
         self.write_pv = write_pv
+        self.converter: CaConverter = DisconnectedCaConverter(float, dbr.DBR_DOUBLE)
         self.initial_values: dict[str, AugmentedValue] = {}
-        self.converter: CaConverter = DisconnectedCaConverter(None, None)
         self.subscription: Subscription | None = None
+        super().__init__(datatype)
 
-    def source(self, name: str):
-        return f"ca://{self.read_pv}"
+    def source(self, name: str, read: bool):
+        return f"ca://{self.read_pv if read else self.write_pv}"
 
-    async def _store_initial_value(self, pv, timeout: float = DEFAULT_TIMEOUT):
+    async def _store_initial_value(self, pv: str, timeout: float):
         try:
             self.initial_values[pv] = await caget(
                 pv, format=FORMAT_CTRL, timeout=timeout
@@ -306,7 +253,7 @@ class CaSignalBackend(SignalBackend[T]):
             logging.debug(f"signal ca://{pv} timed out")
             raise NotConnected(f"ca://{pv}") from exc
 
-    async def connect(self, timeout: float = DEFAULT_TIMEOUT):
+    async def connect(self, timeout: float):
         _use_pyepics_context_if_imported()
         if self.read_pv != self.write_pv:
             # Different, need to connect both
@@ -319,7 +266,19 @@ class CaSignalBackend(SignalBackend[T]):
             await self._store_initial_value(self.read_pv, timeout=timeout)
         self.converter = make_converter(self.datatype, self.initial_values)
 
-    async def put(self, value: T | None, wait=True, timeout=None):
+    async def _caget(self, pv: str, format: Format) -> AugmentedValue:
+        return await caget(
+            pv, datatype=self.converter.read_dbr, format=format, timeout=None
+        )
+
+    def _make_reading(self, value: AugmentedValue) -> Reading[SignalDatatypeT]:
+        return {
+            "value": self.converter.value(value),
+            "timestamp": value.timestamp,
+            "alarm_severity": -1 if value.severity > 2 else value.severity,
+        }
+
+    async def put(self, value: SignalDatatypeT | None, wait: bool):
         if value is None:
             write_value = self.initial_values[self.write_pv]
         else:
@@ -329,50 +288,39 @@ class CaSignalBackend(SignalBackend[T]):
             write_value,
             datatype=self.converter.write_dbr,
             wait=wait,
-            timeout=timeout,
-        )
-
-    async def _caget(self, format: Format) -> AugmentedValue:
-        return await caget(
-            self.read_pv,
-            datatype=self.converter.read_dbr,
-            format=format,
             timeout=None,
         )
 
     async def get_datakey(self, source: str) -> DataKey:
-        value = await self._caget(FORMAT_CTRL)
-        return self.converter.get_datakey(value)
+        value = await self._caget(self.read_pv, FORMAT_CTRL)
+        metadata = _metadata_from_augmented_value(value, self.converter.metadata)
+        return make_datakey(
+            self.converter.datatype, self.converter.value(value), source, metadata
+        )
 
-    async def get_reading(self) -> Reading:
-        value = await self._caget(FORMAT_TIME)
-        return self.converter.reading(value)
+    async def get_reading(self) -> Reading[SignalDatatypeT]:
+        value = await self._caget(self.read_pv, FORMAT_TIME)
+        return self._make_reading(value)
 
-    async def get_value(self) -> T:
-        value = await self._caget(FORMAT_RAW)
+    async def get_value(self) -> SignalDatatypeT:
+        value = await self._caget(self.read_pv, FORMAT_RAW)
         return self.converter.value(value)
 
-    async def get_setpoint(self) -> T:
-        value = await caget(
-            self.write_pv,
-            datatype=self.converter.read_dbr,
-            format=FORMAT_RAW,
-            timeout=None,
-        )
+    async def get_setpoint(self) -> SignalDatatypeT:
+        value = await self._caget(self.write_pv, FORMAT_RAW)
         return self.converter.value(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.subscription
             ), "Cannot set a callback when one is already set"
             self.subscription = camonitor(
                 self.read_pv,
-                lambda v: callback(self.converter.reading(v), self.converter.value(v)),
+                lambda v: callback(self._make_reading(v)),
                 datatype=self.converter.read_dbr,
                 format=FORMAT_TIME,
             )
-        else:
-            if self.subscription:
-                self.subscription.close()
+        elif self.subscription:
+            self.subscription.close()
             self.subscription = None

ophyd_async/epics/signal/_common.py

@@ -1,57 +1,43 @@
-import inspect
-from enum import Enum
+from collections.abc import Sequence
+from typing import Any, get_args, get_origin
 
-from typing_extensions import TypedDict
+import numpy as np
 
-from ophyd_async.core import RuntimeSubsetEnum
-
-common_meta = {
-    "units",
-    "precision",
-}
-
-
-class LimitPair(TypedDict):
-    high: float | None
-    low: float | None
-
-
-class Limits(TypedDict):
-    alarm: LimitPair
-    control: LimitPair
-    display: LimitPair
-    warning: LimitPair
+from ophyd_async.core import SubsetEnum, get_dtype, get_enum_cls
 
 
 def get_supported_values(
     pv: str,
-    datatype: type[str] | None,
-    pv_choices: tuple[str, ...],
+    datatype: type,
+    pv_choices: Sequence[str],
 ) -> dict[str, str]:
-    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."
-            )
-
-        choices = tuple(v.value for v in datatype)
+    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 not set(choices).issubset(pv_choices):
+            raise TypeError(error_msg + "to be a subset of them.")
+    else:
         if set(choices) != set(pv_choices):
-            raise TypeError(
-                f"{pv} has choices {pv_choices}, "
-                f"which do not match {datatype}, which has {choices}."
-            )
-        return {x: datatype(x) if x else "_" for x in pv_choices}
-    elif datatype is None or datatype is str:
-        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."
-    )
+            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
+    return supported_values
+
+
+def format_datatype(datatype: Any) -> str:
+    if get_origin(datatype) is np.ndarray and get_args(datatype)[0] == tuple[int]:
+        dtype = get_dtype(datatype)
+        return f"Array1D[np.{dtype.name}]"
+    elif get_origin(datatype) is Sequence:
+        return f"Sequence[{get_args(datatype)[0].__name__}]"
+    elif isinstance(datatype, type):
+        return datatype.__name__
+    else:
+        return str(datatype)