ophyd-async 0.7.0a1__py3-none-any.whl → 0.8.0a3__py3-none-any.whl

ophyd_async/core/_device_filler.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from collections.abc import Callable, Iterator, Sequence
+from typing import (
+    Any,
+    Generic,
+    NewType,
+    NoReturn,
+    Protocol,
+    TypeVar,
+    cast,
+    get_args,
+    get_type_hints,
+    runtime_checkable,
+)
+
+from ._device import Device, DeviceConnector, DeviceVector
+from ._signal import Signal, SignalX
+from ._signal_backend import SignalBackend, SignalDatatype
+from ._utils import get_origin_class
+
+SignalBackendT = TypeVar("SignalBackendT", bound=SignalBackend)
+DeviceConnectorT = TypeVar("DeviceConnectorT", bound=DeviceConnector)
+# Unique name possibly with trailing understore, the attribute name on the Device
+UniqueName = NewType("UniqueName", str)
+# Logical name without trailing underscore, the name in the control system
+LogicalName = NewType("LogicalName", str)
+
+
+def _get_datatype(annotation: Any) -> type | None:
+    """Return int from SignalRW[int]."""
+    args = get_args(annotation)
+    if len(args) == 1 and get_origin_class(args[0]):
+        return args[0]
+
+
+def _logical(name: UniqueName) -> LogicalName:
+    return LogicalName(name.rstrip("_"))
+
+
+@runtime_checkable
+class DeviceAnnotation(Protocol):
+    @abstractmethod
+    def __call__(self, parent: Device, child: Device): ...
+
+
+class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
+    def __init__(
+        self,
+        device: Device,
+        signal_backend_factory: Callable[[type[SignalDatatype] | None], SignalBackendT],
+        device_connector_factory: Callable[[], DeviceConnectorT],
+    ):
+        self._device = device
+        self._signal_backend_factory = signal_backend_factory
+        self._device_connector_factory = device_connector_factory
+        # Annotations stored ready for the creation phase
+        self._uncreated_signals: dict[UniqueName, type[Signal]] = {}
+        self._uncreated_devices: dict[UniqueName, type[Device]] = {}
+        self._extras: dict[UniqueName, Sequence[Any]] = {}
+        self._signal_datatype: dict[LogicalName, type | None] = {}
+        self._vector_device_type: dict[LogicalName, type[Device] | None] = {}
+        # Backends and Connectors stored ready for the connection phase
+        self._unfilled_backends: dict[
+            LogicalName, tuple[SignalBackendT, type[Signal]]
+        ] = {}
+        self._unfilled_connectors: dict[LogicalName, DeviceConnectorT] = {}
+        # Once they are filled they go here in case we reconnect
+        self._filled_backends: dict[
+            LogicalName, tuple[SignalBackendT, type[Signal]]
+        ] = {}
+        self._filled_connectors: dict[LogicalName, DeviceConnectorT] = {}
+        self._scan_for_annotations()
+
+    def _raise(self, name: str, error: str) -> NoReturn:
+        raise TypeError(f"{type(self._device).__name__}.{name}: {error}")
+
+    def _store_signal_datatype(self, name: UniqueName, annotation: Any):
+        origin = get_origin_class(annotation)
+        datatype = _get_datatype(annotation)
+        if origin == SignalX:
+            # SignalX doesn't need datatype
+            self._signal_datatype[_logical(name)] = None
+        elif origin and issubclass(origin, Signal) and datatype:
+            # All other Signals need one
+            self._signal_datatype[_logical(name)] = datatype
+        else:
+            # Not recognized
+            self._raise(
+                name,
+                f"Expected SignalX or SignalR/W/RW[type], got {annotation}",
+            )
+
+    def _scan_for_annotations(self):
+        # Get type hints on the class, not the instance
+        # https://github.com/python/cpython/issues/124840
+        cls = type(self._device)
+        # Get hints without Annotated for determining types
+        hints = get_type_hints(cls)
+        # Get hints with Annotated for wrapping signals and backends
+        extra_hints = get_type_hints(cls, include_extras=True)
+        for attr_name, annotation in hints.items():
+            name = UniqueName(attr_name)
+            origin = get_origin_class(annotation)
+            if (
+                name == "parent"
+                or name.startswith("_")
+                or not origin
+                or not issubclass(origin, Device)
+            ):
+                # Ignore any child that is not a public Device
+                continue
+            self._extras[name] = getattr(extra_hints[attr_name], "__metadata__", ())
+            if issubclass(origin, Signal):
+                self._store_signal_datatype(name, annotation)
+                self._uncreated_signals[name] = origin
+            elif origin == DeviceVector:
+                child_type = _get_datatype(annotation)
+                child_origin = get_origin_class(child_type)
+                if child_origin is None or not issubclass(child_origin, Device):
+                    self._raise(
+                        name,
+                        f"Expected DeviceVector[SomeDevice], got {annotation}",
+                    )
+                if issubclass(child_origin, Signal):
+                    self._store_signal_datatype(name, child_type)
+                self._vector_device_type[_logical(name)] = child_origin
+                setattr(self._device, name, DeviceVector({}))
+            else:
+                self._uncreated_devices[name] = origin
+
+    def check_created(self):
+        uncreated = sorted(set(self._uncreated_signals).union(self._uncreated_devices))
+        if uncreated:
+            raise RuntimeError(
+                f"{self._device.name}: {uncreated} have not been created yet"
+            )
+
+    def create_signals_from_annotations(
+        self,
+        filled=True,
+    ) -> Iterator[tuple[SignalBackendT, list[Any]]]:
+        for name in list(self._uncreated_signals):
+            child_type = self._uncreated_signals.pop(name)
+            backend = self._signal_backend_factory(
+                self._signal_datatype[_logical(name)]
+            )
+            extras = list(self._extras[name])
+            yield backend, extras
+            signal = child_type(backend)
+            for anno in extras:
+                assert isinstance(anno, DeviceAnnotation), anno
+                anno(self._device, signal)
+            setattr(self._device, name, signal)
+            dest = self._filled_backends if filled else self._unfilled_backends
+            dest[_logical(name)] = (backend, child_type)
+
+    def create_devices_from_annotations(
+        self,
+        filled=True,
+    ) -> Iterator[tuple[DeviceConnectorT, list[Any]]]:
+        for name in list(self._uncreated_devices):
+            child_type = self._uncreated_devices.pop(name)
+            connector = self._device_connector_factory()
+            extras = list(self._extras[name])
+            yield connector, extras
+            device = child_type(connector=connector)
+            for anno in extras:
+                assert isinstance(anno, DeviceAnnotation), anno
+                anno(self._device, device)
+            setattr(self._device, name, device)
+            dest = self._filled_connectors if filled else self._unfilled_connectors
+            dest[_logical(name)] = connector
+
+    def create_device_vector_entries_to_mock(self, num: int):
+        for name, cls in self._vector_device_type.items():
+            assert cls, "Shouldn't happen"
+            for i in range(1, num + 1):
+                if issubclass(cls, Signal):
+                    self.fill_child_signal(name, cls, i)
+                elif issubclass(cls, Device):
+                    self.fill_child_device(name, cls, i)
+                else:
+                    self._raise(name, f"Can't make {cls}")
+
+    def check_filled(self, source: str):
+        unfilled = sorted(set(self._unfilled_connectors).union(self._unfilled_backends))
+        if unfilled:
+            raise RuntimeError(
+                f"{self._device.name}: cannot provision {unfilled} from {source}"
+            )
+
+    def _ensure_device_vector(self, name: LogicalName) -> DeviceVector:
+        if not hasattr(self._device, name):
+            # We have no type hints, so use whatever we are told
+            self._vector_device_type[name] = None
+            setattr(self._device, name, DeviceVector({}))
+        vector = getattr(self._device, name)
+        if not isinstance(vector, DeviceVector):
+            self._raise(name, f"Expected DeviceVector, got {vector}")
+        return vector
+
+    def fill_child_signal(
+        self,
+        name: str,
+        signal_type: type[Signal],
+        vector_index: int | None = None,
+    ) -> SignalBackendT:
+        name = cast(LogicalName, name)
+        if name in self._unfilled_backends:
+            # We made it above
+            backend, expected_signal_type = self._unfilled_backends.pop(name)
+            self._filled_backends[name] = backend, expected_signal_type
+        elif name in self._filled_backends:
+            # We made it and filled it so return for validation
+            backend, expected_signal_type = self._filled_backends[name]
+        elif vector_index:
+            # We need to add a new entry to a DeviceVector
+            vector = self._ensure_device_vector(name)
+            backend = self._signal_backend_factory(self._signal_datatype.get(name))
+            expected_signal_type = self._vector_device_type[name] or signal_type
+            vector[vector_index] = signal_type(backend)
+        elif child := getattr(self._device, name, None):
+            # There is an existing child, so raise
+            self._raise(name, f"Cannot make child as it would shadow {child}")
+        else:
+            # We need to add a new child to the top level Device
+            backend = self._signal_backend_factory(None)
+            expected_signal_type = signal_type
+            setattr(self._device, name, signal_type(backend))
+        if signal_type is not expected_signal_type:
+            self._raise(
+                name,
+                f"is a {signal_type.__name__} not a {expected_signal_type.__name__}",
+            )
+        return backend
+
+    def fill_child_device(
+        self,
+        name: str,
+        device_type: type[Device] = Device,
+        vector_index: int | None = None,
+    ) -> DeviceConnectorT:
+        name = cast(LogicalName, name)
+        if name in self._unfilled_connectors:
+            # We made it above
+            connector = self._unfilled_connectors.pop(name)
+            self._filled_connectors[name] = connector
+        elif name in self._filled_backends:
+            # We made it and filled it so return for validation
+            connector = self._filled_connectors[name]
+        elif vector_index:
+            # We need to add a new entry to a DeviceVector
+            vector = self._ensure_device_vector(name)
+            vector_device_type = self._vector_device_type[name] or device_type
+            assert issubclass(
+                vector_device_type, Device
+            ), f"{vector_device_type} is not a Device"
+            connector = self._device_connector_factory()
+            vector[vector_index] = vector_device_type(connector=connector)
+        elif child := getattr(self._device, name, None):
+            # There is an existing child, so raise
+            self._raise(name, f"Cannot make child as it would shadow {child}")
+        else:
+            # We need to add a new child to the top level Device
+            connector = self._device_connector_factory()
+            setattr(self._device, name, device_type(connector=connector))
+        return connector

ophyd_async/core/_device_save_loader.py

@@ -27,11 +27,8 @@ def pydantic_model_abstraction_representer(
     return dumper.represent_data(model.model_dump(mode="python"))
 
 
-class OphydDumper(yaml.Dumper):
-    def represent_data(self, data: Any) -> Any:
-        if isinstance(data, Enum):
-            return self.represent_data(data.value)
-        return super().represent_data(data)
+def enum_representer(dumper: yaml.Dumper, enum: Enum) -> yaml.Node:
+    return dumper.represent_data(enum.value)
 
 
 def get_signal_values(
@@ -74,7 +71,7 @@ def get_signal_values(
         for key, value in zip(selected_signals, selected_values, strict=False)
     }
     # Ignored values place in with value None so we know which ones were ignored
-    named_values.update({key: None for key in ignore})
+    named_values.update(dict.fromkeys(ignore))
     return named_values
 
 
@@ -111,6 +108,7 @@ def walk_rw_signals(
         path_prefix = ""
 
     signals: dict[str, SignalRW[Any]] = {}
+
     for attr_name, attr in device.children():
         dot_path = f"{path_prefix}{attr_name}"
         if type(attr) is SignalRW:
@@ -145,9 +143,10 @@ def save_to_yaml(phases: Sequence[dict[str, Any]], save_path: str | Path) -> Non
         pydantic_model_abstraction_representer,
         Dumper=yaml.Dumper,
     )
+    yaml.add_multi_representer(Enum, enum_representer, Dumper=yaml.Dumper)
 
     with open(save_path, "w") as file:
-        yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)
+        yaml.dump(phases, file)
 
 
 def load_from_yaml(save_path: str) -> Sequence[dict[str, Any]]:

ophyd_async/core/_mock_signal_backend.py

@@ -1,84 +1,76 @@
 import asyncio
 from collections.abc import Callable
 from functools import cached_property
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock
 
 from bluesky.protocols import Descriptor, Reading
 
-from ._signal_backend import SignalBackend
+from ._signal_backend import SignalBackend, SignalDatatypeT
 from ._soft_signal_backend import SoftSignalBackend
-from ._utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
+from ._utils import Callback
 
 
-class MockSignalBackend(SignalBackend[T]):
+class MockSignalBackend(SignalBackend[SignalDatatypeT]):
     """Signal backend for testing, created by ``Device.connect(mock=True)``."""
 
     def __init__(
         self,
-        datatype: type[T] | None = None,
-        initial_backend: SignalBackend[T] | None = None,
+        initial_backend: SignalBackend[SignalDatatypeT],
+        mock: Mock,
     ) -> None:
         if isinstance(initial_backend, MockSignalBackend):
-            raise ValueError("Cannot make a MockSignalBackend for a MockSignalBackends")
+            raise ValueError("Cannot make a MockSignalBackend for a MockSignalBackend")
 
         self.initial_backend = initial_backend
 
-        if datatype is None:
-            assert (
-                self.initial_backend
-            ), "Must supply either initial_backend or datatype"
-            datatype = self.initial_backend.datatype
+        if isinstance(self.initial_backend, SoftSignalBackend):
+            # Backend is already a SoftSignalBackend, so use it
+            self.soft_backend = self.initial_backend
+        else:
+            # Backend is not a SoftSignalBackend, so create one to mimic it
+            self.soft_backend = SoftSignalBackend(
+                datatype=self.initial_backend.datatype
+            )
 
-        self.datatype = datatype
+        # use existing Mock if provided
+        self.mock = mock
+        self.put_mock = AsyncMock(name="put", spec=Callable)
+        self.mock.attach_mock(self.put_mock, "put")
 
-        if not isinstance(self.initial_backend, SoftSignalBackend):
-            # If the backend is a hard signal backend, or not provided,
-            # then we create a soft signal to mimic it
+        super().__init__(datatype=self.initial_backend.datatype)
 
-            self.soft_backend = SoftSignalBackend(datatype=datatype)
-        else:
-            self.soft_backend = self.initial_backend
+    def set_value(self, value: SignalDatatypeT):
+        self.soft_backend.set_value(value)
 
-    def source(self, name: str) -> str:
-        if self.initial_backend:
-            return f"mock+{self.initial_backend.source(name)}"
-        return f"mock+{name}"
+    def source(self, name: str, read: bool) -> str:
+        return f"mock+{self.initial_backend.source(name, read)}"
 
-    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
+    async def connect(self, timeout: float) -> None:
         pass
 
-    @cached_property
-    def put_mock(self) -> AsyncMock:
-        return AsyncMock(name="put", spec=Callable)
-
     @cached_property
     def put_proceeds(self) -> asyncio.Event:
         put_proceeds = asyncio.Event()
         put_proceeds.set()
         return put_proceeds
 
-    async def put(self, value: T | None, wait=True, timeout=None):
-        await self.put_mock(value, wait=wait, timeout=timeout)
-        await self.soft_backend.put(value, wait=wait, timeout=timeout)
-
+    async def put(self, value: SignalDatatypeT | None, wait: bool):
+        await self.put_mock(value, wait=wait)
+        await self.soft_backend.put(value, wait=wait)
         if wait:
-            await asyncio.wait_for(self.put_proceeds.wait(), timeout=timeout)
-
-    def set_value(self, value: T):
-        self.soft_backend.set_value(value)
+            await self.put_proceeds.wait()
 
     async def get_reading(self) -> Reading:
         return await self.soft_backend.get_reading()
 
-    async def get_value(self) -> T:
+    async def get_value(self) -> SignalDatatypeT:
         return await self.soft_backend.get_value()
 
-    async def get_setpoint(self) -> T:
-        """For a soft signal, the setpoint and readback values are the same."""
+    async def get_setpoint(self) -> SignalDatatypeT:
         return await self.soft_backend.get_setpoint()
 
     async def get_datakey(self, source: str) -> Descriptor:
         return await self.soft_backend.get_datakey(source)
 
-    def set_callback(self, callback: ReadingValueCallback[T] | None) -> None:
+    def set_callback(self, callback: Callback[Reading[SignalDatatypeT]] | None) -> None:
         self.soft_backend.set_callback(callback)

ophyd_async/core/_mock_signal_utils.py

@@ -1,23 +1,21 @@
 from collections.abc import Awaitable, Callable, Iterable
 from contextlib import asynccontextmanager, contextmanager
-from typing import Any
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock
 
+from ._device import Device, _device_mocks
 from ._mock_signal_backend import MockSignalBackend
-from ._signal import Signal
-from ._utils import T
+from ._signal import Signal, SignalR, _mock_signal_backends
+from ._soft_signal_backend import SignalDatatypeT
 
 
 def _get_mock_signal_backend(signal: Signal) -> MockSignalBackend:
-    backend = signal._backend  # noqa:SLF001
-    assert isinstance(backend, MockSignalBackend), (
-        "Expected to receive a `MockSignalBackend`, instead "
-        f" received {type(backend)}. "
-    )
-    return backend
+    assert (
+        signal in _mock_signal_backends
+    ), f"Signal {signal} not connected in mock mode"
+    return _mock_signal_backends[signal]
 
 
-def set_mock_value(signal: Signal[T], value: T):
+def set_mock_value(signal: Signal[SignalDatatypeT], value: SignalDatatypeT):
     """Set the value of a signal that is in mock mode."""
     backend = _get_mock_signal_backend(signal)
     backend.set_value(value)
@@ -47,6 +45,12 @@ def get_mock_put(signal: Signal) -> AsyncMock:
     return _get_mock_signal_backend(signal).put_mock
 
 
+def get_mock(device: Device | Signal) -> Mock:
+    if isinstance(device, Signal):
+        return _get_mock_signal_backend(device).mock
+    return _device_mocks[device]
+
+
 def reset_mock_put_calls(signal: Signal):
     backend = _get_mock_signal_backend(signal)
     backend.put_mock.reset_mock()
@@ -59,8 +63,8 @@ class _SetValuesIterator:
 
     def __init__(
         self,
-        signal: Signal,
-        values: Iterable[Any],
+        signal: SignalR[SignalDatatypeT],
+        values: Iterable[SignalDatatypeT],
         require_all_consumed: bool = False,
     ):
         self.signal = signal
@@ -99,8 +103,8 @@ class _SetValuesIterator:
 
 
 def set_mock_values(
-    signal: Signal,
-    values: Iterable[Any],
+    signal: SignalR[SignalDatatypeT],
+    values: Iterable[SignalDatatypeT],
     require_all_consumed: bool = False,
 ) -> _SetValuesIterator:
     """Iterator to set a signal to a sequence of values, optionally repeating the
@@ -143,7 +147,9 @@ def _unset_side_effect_cm(put_mock: AsyncMock):
 
 
 def callback_on_mock_put(
-    signal: Signal[T], callback: Callable[[T], None] | Callable[[T], Awaitable[None]]
+    signal: Signal[SignalDatatypeT],
+    callback: Callable[[SignalDatatypeT, bool], None]
+    | Callable[[SignalDatatypeT, bool], Awaitable[None]],
 ):
     """For setting a callback when a backend is put to.
 

ophyd_async/core/_protocol.py

@@ -13,10 +13,38 @@ from typing import (
 from bluesky.protocols import HasName, Reading
 from event_model import DataKey
 
+from ._utils import DEFAULT_TIMEOUT
+
 if TYPE_CHECKING:
+    from unittest.mock import Mock
+
     from ._status import AsyncStatus
 
 
+@runtime_checkable
+class Connectable(Protocol):
+    @abstractmethod
+    async def connect(
+        self,
+        mock: bool | Mock = False,
+        timeout: float = DEFAULT_TIMEOUT,
+        force_reconnect: bool = False,
+    ):
+        """Connect self and all child Devices.
+
+        Contains a timeout that gets propagated to child.connect methods.
+
+        Parameters
+        ----------
+        mock:
+            If True then use ``MockSignalBackend`` for all Signals
+        timeout:
+            Time to wait before failing with a TimeoutError.
+        force_reconnect:
+            Reconnect even if previous connect was successful.
+        """
+
+
 @runtime_checkable
 class AsyncReadable(HasName, Protocol):
     @abstractmethod
@@ -33,7 +61,6 @@ class AsyncReadable(HasName, Protocol):
                          ('channel2',
                          {'value': 16, 'timestamp': 1472493713.539238}))
         """
-        ...
 
     @abstractmethod
     async def describe(self) -> dict[str, DataKey]:
@@ -53,7 +80,6 @@ class AsyncReadable(HasName, Protocol):
                           'dtype': 'number',
                           'shape': []}))
         """
-        ...
 
 
 @runtime_checkable
@@ -63,14 +89,12 @@ class AsyncConfigurable(HasName, Protocol):
         """Same API as ``read`` but for slow-changing fields related to configuration.
         e.g., exposure time. These will typically be read only once per run.
         """
-        ...
 
     @abstractmethod
     async def describe_configuration(self) -> dict[str, DataKey]:
         """Same API as ``describe``, but corresponding to the keys in
         ``read_configuration``.
         """
-        ...
 
 
 @runtime_checkable
@@ -78,12 +102,10 @@ class AsyncPausable(Protocol):
     @abstractmethod
     async def pause(self) -> None:
         """Perform device-specific work when the RunEngine pauses."""
-        ...
 
     @abstractmethod
     async def resume(self) -> None:
         """Perform device-specific work when the RunEngine resumes after a pause."""
-        ...
 
 
 @runtime_checkable
@@ -95,7 +117,6 @@ class AsyncStageable(Protocol):
         It should return a ``Status`` that is marked done when the device is
         done staging.
         """
-        ...
 
     @abstractmethod
     def unstage(self) -> AsyncStatus:
@@ -104,7 +125,6 @@ class AsyncStageable(Protocol):
         It should return a ``Status`` that is marked done when the device is finished
         unstaging.
         """
-        ...
 
 
 C = TypeVar("C", contravariant=True)