ophyd-async 0.3a3__py3-none-any.whl → 0.3a5__py3-none-any.whl

ophyd_async/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3a3'
+__version__ = version = '0.3a5'
 __version_tuple__ = version_tuple = (0, 3)

ophyd_async/core/__init__.py

@@ -5,7 +5,7 @@ from ._providers import (
     ShapeProvider,
     StaticDirectoryProvider,
 )
-from .async_status import AsyncStatus
+from .async_status import AsyncStatus, WatchableAsyncStatus
 from .detector import (
     DetectorControl,
     DetectorTrigger,
@@ -24,6 +24,18 @@ from .device_save_loader import (
     walk_rw_signals,
 )
 from .flyer import HardwareTriggeredFlyable, TriggerLogic
+from .mock_signal_backend import (
+    MockSignalBackend,
+)
+from .mock_signal_utils import (
+    assert_mock_put_called_with,
+    callback_on_mock_put,
+    mock_puts_blocked,
+    reset_mock_put_calls,
+    set_mock_put_proceeds,
+    set_mock_value,
+    set_mock_values,
+)
 from .signal import (
     Signal,
     SignalR,
@@ -36,15 +48,12 @@ from .signal import (
     assert_value,
     observe_value,
     set_and_wait_for_value,
-    set_sim_callback,
-    set_sim_put_proceeds,
-    set_sim_value,
-    soft_signal_r_and_backend,
+    soft_signal_r_and_setter,
     soft_signal_rw,
     wait_for_value,
 )
 from .signal_backend import SignalBackend
-from .sim_signal_backend import SimSignalBackend
+from .soft_signal_backend import SoftSignalBackend
 from .standard_readable import ConfigSignal, HintedSignal, StandardReadable
 from .utils import (
     DEFAULT_TIMEOUT,
@@ -59,9 +68,15 @@ from .utils import (
 )
 
 __all__ = [
+    "assert_mock_put_called_with",
+    "callback_on_mock_put",
+    "mock_puts_blocked",
+    "set_mock_values",
+    "reset_mock_put_calls",
     "SignalBackend",
-    "SimSignalBackend",
+    "SoftSignalBackend",
     "DetectorControl",
+    "MockSignalBackend",
     "DetectorTrigger",
     "DetectorWriter",
     "StandardDetector",
@@ -73,15 +88,15 @@ __all__ = [
     "SignalW",
     "SignalRW",
     "SignalX",
-    "soft_signal_r_and_backend",
+    "soft_signal_r_and_setter",
     "soft_signal_rw",
     "observe_value",
     "set_and_wait_for_value",
-    "set_sim_callback",
-    "set_sim_put_proceeds",
-    "set_sim_value",
+    "set_mock_put_proceeds",
+    "set_mock_value",
     "wait_for_value",
     "AsyncStatus",
+    "WatchableAsyncStatus",
     "DirectoryInfo",
     "DirectoryProvider",
     "NameProvider",

ophyd_async/core/async_status.py

@@ -2,28 +2,42 @@
 
 import asyncio
 import functools
-from typing import Awaitable, Callable, Coroutine, List, Optional, cast
+import time
+from dataclasses import asdict, replace
+from typing import (
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Generic,
+    SupportsFloat,
+    Type,
+    TypeVar,
+    cast,
+)
 
 from bluesky.protocols import Status
 
-from .utils import Callback, T
+from ..protocols import Watcher
+from .utils import Callback, P, T, WatcherUpdate
 
+AS = TypeVar("AS", bound="AsyncStatus")
+WAS = TypeVar("WAS", bound="WatchableAsyncStatus")
 
-class AsyncStatus(Status):
+
+class AsyncStatusBase(Status):
     """Convert asyncio awaitable to bluesky Status interface"""
 
-    def __init__(
-        self,
-        awaitable: Awaitable,
-        watchers: Optional[List[Callable]] = None,
-    ):
+    def __init__(self, awaitable: Awaitable, timeout: SupportsFloat | None = None):
+        if isinstance(timeout, SupportsFloat):
+            timeout = float(timeout)
         if isinstance(awaitable, asyncio.Task):
             self.task = awaitable
         else:
-            self.task = asyncio.create_task(awaitable)  # type: ignore
+            self.task = asyncio.create_task(
+                asyncio.wait_for(awaitable, timeout=timeout)
+            )
         self.task.add_done_callback(self._run_callbacks)
-        self._callbacks = cast(List[Callback[Status]], [])
-        self._watchers = watchers
+        self._callbacks: list[Callback[Status]] = []
 
     def __await__(self):
         return self.task.__await__()
@@ -39,15 +53,11 @@ class AsyncStatus(Status):
             for callback in self._callbacks:
                 callback(self)
 
-    # TODO: remove ignore and bump min version when bluesky v1.12.0 is released
-    def exception(  # type: ignore
-        self, timeout: Optional[float] = 0.0
-    ) -> Optional[BaseException]:
+    def exception(self, timeout: float | None = 0.0) -> BaseException | None:
         if timeout != 0.0:
-            raise Exception(
+            raise ValueError(
                 "cannot honour any timeout other than 0 in an asynchronous function"
             )
-
         if self.task.done():
             try:
                 return self.task.exception()
@@ -67,22 +77,6 @@ class AsyncStatus(Status):
             and self.task.exception() is None
         )
 
-    def watch(self, watcher: Callable):
-        """Add watcher to the list of interested parties.
-
-        Arguments as per Bluesky :external+bluesky:meth:`watch` protocol.
-        """
-        if self._watchers is not None:
-            self._watchers.append(watcher)
-
-    @classmethod
-    def wrap(cls, f: Callable[[T], Coroutine]) -> Callable[[T], "AsyncStatus"]:
-        @functools.wraps(f)
-        def wrap_f(self) -> AsyncStatus:
-            return AsyncStatus(f(self))
-
-        return wrap_f
-
     def __repr__(self) -> str:
         if self.done:
             if e := self.exception():
@@ -94,3 +88,72 @@ class AsyncStatus(Status):
         return f"<{type(self).__name__}, task: {self.task.get_coro()}, {status}>"
 
     __str__ = __repr__
+
+
+class AsyncStatus(AsyncStatusBase):
+    @classmethod
+    def wrap(cls: Type[AS], f: Callable[P, Awaitable]) -> Callable[P, AS]:
+        @functools.wraps(f)
+        def wrap_f(*args: P.args, **kwargs: P.kwargs) -> AS:
+            # We can't type this more properly because Concatenate/ParamSpec doesn't
+            # yet support keywords
+            # https://peps.python.org/pep-0612/#concatenating-keyword-parameters
+            timeout = kwargs.get("timeout")
+            assert isinstance(timeout, SupportsFloat) or timeout is None
+            return cls(f(*args, **kwargs), timeout=timeout)
+
+        # type is actually functools._Wrapped[P, Awaitable, P, AS]
+        # but functools._Wrapped is not necessarily available
+        return cast(Callable[P, AS], wrap_f)
+
+
+class WatchableAsyncStatus(AsyncStatusBase, Generic[T]):
+    """Convert AsyncIterator of WatcherUpdates to bluesky Status interface."""
+
+    def __init__(
+        self,
+        iterator: AsyncIterator[WatcherUpdate[T]],
+        timeout: SupportsFloat | None = None,
+    ):
+        self._watchers: list[Watcher] = []
+        self._start = time.monotonic()
+        self._last_update: WatcherUpdate[T] | None = None
+        super().__init__(self._notify_watchers_from(iterator), timeout)
+
+    async def _notify_watchers_from(self, iterator: AsyncIterator[WatcherUpdate[T]]):
+        async for update in iterator:
+            self._last_update = (
+                update
+                if update.time_elapsed is not None
+                else replace(update, time_elapsed=time.monotonic() - self._start)
+            )
+            for watcher in self._watchers:
+                self._update_watcher(watcher, self._last_update)
+
+    def _update_watcher(self, watcher: Watcher, update: WatcherUpdate[T]):
+        vals = asdict(
+            update, dict_factory=lambda d: {k: v for k, v in d if v is not None}
+        )
+        watcher(**vals)
+
+    def watch(self, watcher: Watcher):
+        self._watchers.append(watcher)
+        if self._last_update:
+            self._update_watcher(watcher, self._last_update)
+
+    @classmethod
+    def wrap(
+        cls: Type[WAS],
+        f: Callable[P, AsyncIterator[WatcherUpdate[T]]],
+    ) -> Callable[P, WAS]:
+        """Wrap an AsyncIterator in a WatchableAsyncStatus. If it takes
+        'timeout' as an argument, this must be a float or None, and it
+        will be propagated to the status."""
+
+        @functools.wraps(f)
+        def wrap_f(*args: P.args, **kwargs: P.kwargs) -> WAS:
+            timeout = kwargs.get("timeout")
+            assert isinstance(timeout, SupportsFloat) or timeout is None
+            return cls(f(*args, **kwargs), timeout=timeout)
+
+        return cast(Callable[P, WAS], wrap_f)

ophyd_async/core/detector.py

@@ -31,9 +31,9 @@ from bluesky.protocols import (
 
 from ophyd_async.protocols import AsyncConfigurable, AsyncReadable
 
-from .async_status import AsyncStatus
+from .async_status import AsyncStatus, WatchableAsyncStatus
 from .device import Device
-from .utils import DEFAULT_TIMEOUT, merge_gathered_dicts
+from .utils import DEFAULT_TIMEOUT, WatcherUpdate, merge_gathered_dicts
 
 T = TypeVar("T")
 
@@ -188,7 +188,7 @@ class StandardDetector(
         self._trigger_info: Optional[TriggerInfo] = None
         # For kickoff
         self._watchers: List[Callable] = []
-        self._fly_status: Optional[AsyncStatus] = None
+        self._fly_status: Optional[WatchableAsyncStatus] = None
         self._fly_start: float
 
         self._intial_frame: int
@@ -292,43 +292,37 @@ class StandardDetector(
             f"Detector {self.controller} needs at least {required}s deadtime, "
             f"but trigger logic provides only {self._trigger_info.deadtime}s"
         )
-
         self._arm_status = await self.controller.arm(
             num=self._trigger_info.num,
             trigger=self._trigger_info.trigger,
             exposure=self._trigger_info.livetime,
         )
-
-    @AsyncStatus.wrap
-    async def kickoff(self) -> None:
-        self._fly_status = AsyncStatus(self._fly(), self._watchers)
         self._fly_start = time.monotonic()
 
-    async def _fly(self) -> None:
-        await self._observe_writer_indicies(self._last_frame)
-
-    async def _observe_writer_indicies(self, end_observation: int):
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        if not self._arm_status:
+            raise Exception("Detector not armed!")
+
+    @WatchableAsyncStatus.wrap
+    async def complete(self):
+        assert self._arm_status, "Prepare not run"
+        assert self._trigger_info
         async for index in self.writer.observe_indices_written(
             self._frame_writing_timeout
         ):
-            for watcher in self._watchers:
-                watcher(
-                    name=self.name,
-                    current=index,
-                    initial=self._initial_frame,
-                    target=end_observation,
-                    unit="",
-                    precision=0,
-                    time_elapsed=time.monotonic() - self._fly_start,
-                )
-            if index >= end_observation:
+            yield WatcherUpdate(
+                name=self.name,
+                current=index,
+                initial=self._initial_frame,
+                target=self._trigger_info.num,
+                unit="",
+                precision=0,
+                time_elapsed=time.monotonic() - self._fly_start,
+            )
+            if index >= self._trigger_info.num:
                 break
 
-    @AsyncStatus.wrap
-    async def complete(self) -> AsyncStatus:
-        assert self._fly_status, "Kickoff not run"
-        return await self._fly_status
-
     async def describe_collect(self) -> Dict[str, DataKey]:
         return self._describe
 

ophyd_async/core/device.py

@@ -2,7 +2,10 @@
 
 from __future__ import annotations
 
+import asyncio
 import sys
+from functools import cached_property
+from logging import LoggerAdapter, getLogger
 from typing import (
     Any,
     Coroutine,
@@ -30,6 +33,9 @@ class Device(HasName):
     _name: str = ""
     #: The parent Device if it exists
     parent: Optional[Device] = None
+    # None if connect hasn't started, a Task if it has
+    _connect_task: Optional[asyncio.Task] = None
+    _connect_mock_arg: bool = False
 
     def __init__(self, name: str = "") -> None:
         self.set_name(name)
@@ -39,6 +45,12 @@ class Device(HasName):
         """Return the name of the Device"""
         return self._name
 
+    @cached_property
+    def log(self):
+        return LoggerAdapter(
+            getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
+        )
+
     def children(self) -> Iterator[Tuple[str, Device]]:
         for attr_name, attr in self.__dict__.items():
             if attr_name != "parent" and isinstance(attr, Device):
@@ -52,30 +64,54 @@ class Device(HasName):
         name:
             New name to set
         """
+
+        # Ensure self.log is recreated after a name change
+        if hasattr(self, "log"):
+            del self.log
+
         self._name = name
         for attr_name, child in self.children():
             child_name = f"{name}-{attr_name.rstrip('_')}" if name else ""
             child.set_name(child_name)
             child.parent = self
 
-    async def connect(self, sim: bool = False, timeout: float = DEFAULT_TIMEOUT):
+    async def connect(
+        self,
+        mock: bool = 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
         ----------
-        sim:
-            If True then connect in simulation mode.
+        mock:
+            If True then use ``MockSignalBackend`` for all Signals
         timeout:
             Time to wait before failing with a TimeoutError.
         """
-        coros = {
-            name: child_device.connect(sim, timeout=timeout)
-            for name, child_device in self.children()
-        }
-        if coros:
-            await wait_for_connection(**coros)
+        # If previous connect with same args has started and not errored, can use it
+        can_use_previous_connect = (
+            self._connect_task
+            and not (self._connect_task.done() and self._connect_task.exception())
+            and self._connect_mock_arg == mock
+        )
+        if force_reconnect or not can_use_previous_connect:
+            # Kick off a connection
+            coros = {
+                name: child_device.connect(
+                    mock, timeout=timeout, force_reconnect=force_reconnect
+                )
+                for name, child_device in self.children()
+            }
+            self._connect_task = asyncio.create_task(wait_for_connection(**coros))
+            self._connect_mock_arg = mock
+
+        assert self._connect_task, "Connect task not created, this shouldn't happen"
+        # Wait for it to complete
+        await self._connect_task
 
 
 VT = TypeVar("VT", bound=Device)
@@ -105,9 +141,9 @@ class DeviceCollector:
         If True, call ``device.set_name(variable_name)`` on all collected
         Devices
     connect:
-        If True, call ``device.connect(sim)`` in parallel on all
+        If True, call ``device.connect(mock)`` in parallel on all
         collected Devices
-    sim:
+    mock:
         If True, connect Signals in simulation mode
     timeout:
         How long to wait for connect before logging an exception
@@ -129,12 +165,12 @@ class DeviceCollector:
         self,
         set_name=True,
         connect=True,
-        sim=False,
+        mock=False,
         timeout: float = 10.0,
     ):
         self._set_name = set_name
         self._connect = connect
-        self._sim = sim
+        self._mock = mock
         self._timeout = timeout
         self._names_on_enter: Set[str] = set()
         self._objects_on_exit: Dict[str, Any] = {}
@@ -168,7 +204,7 @@ class DeviceCollector:
                     obj.set_name(name)
                 if self._connect:
                     connect_coroutines[name] = obj.connect(
-                        self._sim, timeout=self._timeout
+                        self._mock, timeout=self._timeout
                     )
 
         # Connect to all the devices

ophyd_async/core/mock_signal_backend.py

@@ -0,0 +1,85 @@
+import asyncio
+from functools import cached_property
+from typing import Optional, Type
+from unittest.mock import Mock
+
+from bluesky.protocols import Descriptor, Reading
+
+from ophyd_async.core.signal_backend import SignalBackend
+from ophyd_async.core.soft_signal_backend import SoftSignalBackend
+from ophyd_async.core.utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
+
+
+class MockSignalBackend(SignalBackend):
+    def __init__(
+        self,
+        datatype: Optional[Type[T]] = None,
+        initial_backend: Optional[SignalBackend[T]] = None,
+    ) -> None:
+        if isinstance(initial_backend, MockSignalBackend):
+            raise ValueError("Cannot make a MockSignalBackend for a MockSignalBackends")
+
+        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
+
+        self.datatype = datatype
+
+        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 mimick it
+
+            self.soft_backend = SoftSignalBackend(datatype=datatype)
+        else:
+            self.soft_backend = initial_backend
+
+    def source(self, name: str) -> str:
+        if self.initial_backend:
+            return f"mock+{self.initial_backend.source(name)}"
+        return f"mock+{name}"
+
+    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
+        pass
+
+    @cached_property
+    def put_mock(self) -> Mock:
+        return Mock(name="put")
+
+    @cached_property
+    def put_proceeds(self) -> asyncio.Event:
+        put_proceeds = asyncio.Event()
+        put_proceeds.set()
+        return put_proceeds
+
+    async def put(self, value: Optional[T], wait=True, timeout=None):
+        self.put_mock(value, wait=wait, timeout=timeout)
+        await self.soft_backend.put(value, wait=wait, timeout=timeout)
+
+        if wait:
+            await asyncio.wait_for(self.put_proceeds.wait(), timeout=timeout)
+
+    def set_value(self, value: T):
+        self.soft_backend.set_value(value)
+
+    async def get_descriptor(self, source: str) -> Descriptor:
+        return await self.soft_backend.get_descriptor(source)
+
+    async def get_reading(self) -> Reading:
+        return await self.soft_backend.get_reading()
+
+    async def get_value(self) -> T:
+        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."""
+        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: Optional[ReadingValueCallback[T]]) -> None:
+        self.soft_backend.set_callback(callback)