ophyd-async 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

ophyd_async/core/detector.py

@@ -1,27 +1,46 @@
 """Module which defines abstract classes to work with detectors"""
+
 import asyncio
+import time
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from enum import Enum
-from typing import AsyncIterator, Dict, Optional, Sequence
+from typing import (
+    AsyncGenerator,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Sequence,
+    TypeVar,
+)
 
 from bluesky.protocols import (
-    Asset,
-    Configurable,
-    Descriptor,
-    Readable,
+    Collectable,
+    DataKey,
+    Flyable,
+    Preparable,
     Reading,
     Stageable,
+    StreamAsset,
     Triggerable,
-    WritesExternalAssets,
+    WritesStreamAssets,
 )
 
-from .async_status import AsyncStatus
+from ophyd_async.protocols import AsyncConfigurable, AsyncReadable
+
+from .async_status import AsyncStatus, WatchableAsyncStatus
 from .device import Device
-from .signal import SignalR
-from .utils import DEFAULT_TIMEOUT, merge_gathered_dicts
+from .utils import DEFAULT_TIMEOUT, WatcherUpdate, merge_gathered_dicts
+
+T = TypeVar("T")
 
 
 class DetectorTrigger(str, Enum):
+    """Type of mechanism for triggering a detector to take frames"""
+
     #: Detector generates internal trigger for given rate
     internal = "internal"
     #: Expect a series of arbitrary length trigger signals
@@ -32,7 +51,26 @@ class DetectorTrigger(str, Enum):
     variable_gate = "variable_gate"
 
 
+@dataclass(frozen=True)
+class TriggerInfo:
+    """Minimal set of information required to setup triggering on a detector"""
+
+    #: Number of triggers that will be sent
+    num: int
+    #: Sort of triggers that will be sent
+    trigger: DetectorTrigger
+    #: What is the minimum deadtime between triggers
+    deadtime: float
+    #: What is the maximum high time of the triggers
+    livetime: float
+
+
 class DetectorControl(ABC):
+    """
+    Classes implementing this interface should hold the logic for
+    arming and disarming a detector
+    """
+
     @abstractmethod
     def get_deadtime(self, exposure: float) -> float:
         """For a given exposure, how long should the time between exposures be"""
@@ -40,23 +78,38 @@ class DetectorControl(ABC):
     @abstractmethod
     async def arm(
         self,
+        num: int,
         trigger: DetectorTrigger = DetectorTrigger.internal,
-        num: int = 0,
         exposure: Optional[float] = None,
     ) -> AsyncStatus:
-        """Arm the detector and return AsyncStatus.
+        """
+        Arm detector, do all necessary steps to prepare detector for triggers.
+
+        Args:
+            num: Expected number of frames
+            trigger: Type of trigger for which to prepare the detector. Defaults to
+            DetectorTrigger.internal.
+            exposure: Exposure time with which to set up the detector. Defaults to None
+            if not applicable or the detector is expected to use its previously-set
+            exposure time.
 
-        Awaiting the return value will wait for num frames to be written.
+        Returns:
+            AsyncStatus: Status representing the arm operation. This function returning
+            represents the start of the arm. The returned status completing means
+            the detector is now armed.
         """
 
     @abstractmethod
     async def disarm(self):
-        """Disarm the detector"""
+        """Disarm the detector, return detector to an idle state"""
 
 
 class DetectorWriter(ABC):
+    """Logic for making a detector write data to somewhere persistent
+    (e.g. an HDF5 file)"""
+
     @abstractmethod
-    async def open(self, multiplier: int = 1) -> Dict[str, Descriptor]:
+    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
         """Open writer and wait for it to be ready for data.
 
         Args:
@@ -68,65 +121,78 @@ class DetectorWriter(ABC):
         """
 
     @abstractmethod
-    async def wait_for_index(
-        self, index: int, timeout: Optional[float] = DEFAULT_TIMEOUT
-    ) -> None:
-        """Wait until a specific index is ready to be collected"""
+    def observe_indices_written(
+        self, timeout=DEFAULT_TIMEOUT
+    ) -> AsyncGenerator[int, None]:
+        """Yield the index of each frame (or equivalent data point) as it is written"""
 
     @abstractmethod
     async def get_indices_written(self) -> int:
         """Get the number of indices written"""
 
     @abstractmethod
-    def collect_stream_docs(self, indices_written: int) -> AsyncIterator[Asset]:
+    def collect_stream_docs(self, indices_written: int) -> AsyncIterator[StreamAsset]:
         """Create Stream docs up to given number written"""
 
     @abstractmethod
     async def close(self) -> None:
-        """Close writer and wait for it to be finished"""
+        """Close writer, blocks until I/O is complete"""
 
 
 class StandardDetector(
     Device,
     Stageable,
-    Configurable,
-    Readable,
+    AsyncConfigurable,
+    AsyncReadable,
     Triggerable,
-    WritesExternalAssets,
+    Preparable,
+    Flyable,
+    Collectable,
+    WritesStreamAssets,
+    Generic[T],
 ):
-    """Detector with useful default behaviour.
-
-    Must be supplied instances of classes that inherit from DetectorControl and
-    DetectorData, to dictate how the detector will be controlled (i.e. arming and
-    disarming) as well as how the detector data will be written (i.e. opening and
-    closing the writer, and handling data writing indices).
-
-    NOTE: only for step-scans.
+    """
+    Useful detector base class for step and fly scanning detectors.
+    Aggregates controller and writer logic together.
     """
 
     def __init__(
         self,
         controller: DetectorControl,
         writer: DetectorWriter,
-        config_sigs: Sequence[SignalR] = (),
+        config_sigs: Sequence[AsyncReadable] = (),
         name: str = "",
         writer_timeout: float = DEFAULT_TIMEOUT,
     ) -> None:
         """
-        Parameters
-        ----------
-        control:
-            instance of class which inherits from :class:`DetectorControl`
-        data:
-            instance of class which inherits from :class:`DetectorData`
-        name:
-            detector name
+        Constructor
+
+        Args:
+            controller: Logic for arming and disarming the detector
+            writer: Logic for making the detector write persistent data
+            config_sigs: Signals to read when describe and read
+            configuration are called. Defaults to ().
+            name: Device name. Defaults to "".
+            writer_timeout: Timeout for frame writing to start, if the
+            timeout is reached, ophyd-async assumes the detector
+            has a problem and raises an error.
+            Defaults to DEFAULT_TIMEOUT.
         """
         self._controller = controller
         self._writer = writer
-        self._describe: Dict[str, Descriptor] = {}
+        self._describe: Dict[str, DataKey] = {}
         self._config_sigs = list(config_sigs)
         self._frame_writing_timeout = writer_timeout
+        # For prepare
+        self._arm_status: Optional[AsyncStatus] = None
+        self._trigger_info: Optional[TriggerInfo] = None
+        # For kickoff
+        self._watchers: List[Callable] = []
+        self._fly_status: Optional[WatchableAsyncStatus] = None
+        self._fly_start: float
+
+        self._intial_frame: int
+        self._last_frame: int
         super().__init__(name)
 
     @property
@@ -137,14 +203,20 @@ class StandardDetector(
     def writer(self) -> DetectorWriter:
         return self._writer
 
-    async def check_config_sigs(self):
+    @AsyncStatus.wrap
+    async def stage(self) -> None:
+        # Disarm the detector, stop filewriting, and open file for writing.
+        await self._check_config_sigs()
+        await asyncio.gather(self.writer.close(), self.controller.disarm())
+        self._describe = await self.writer.open()
+
+    async def _check_config_sigs(self):
         """Checks configuration signals are named and connected."""
         for signal in self._config_sigs:
-            if signal._name == "":
+            if signal.name == "":
                 raise Exception(
                     "config signal must be named before it is passed to the detector"
                 )
-
             try:
                 await signal.get_value()
             except NotImplementedError:
@@ -154,46 +226,116 @@ class StandardDetector(
                 )
 
     @AsyncStatus.wrap
-    async def stage(self) -> None:
-        """Disarm the detector, stop filewriting, and open file for writing."""
-        await self.check_config_sigs()
-        await asyncio.gather(self.writer.close(), self.controller.disarm())
-        self._describe = await self.writer.open()
-
-    async def describe_configuration(self) -> Dict[str, Descriptor]:
-        return await merge_gathered_dicts(sig.describe() for sig in self._config_sigs)
+    async def unstage(self) -> None:
+        # Stop data writing.
+        await self.writer.close()
 
     async def read_configuration(self) -> Dict[str, Reading]:
         return await merge_gathered_dicts(sig.read() for sig in self._config_sigs)
 
-    def describe(self) -> Dict[str, Descriptor]:
+    async def describe_configuration(self) -> Dict[str, DataKey]:
+        return await merge_gathered_dicts(sig.describe() for sig in self._config_sigs)
+
+    async def read(self) -> Dict[str, Reading]:
+        # All data is in StreamResources, not Events, so nothing to output here
+        return {}
+
+    async def describe(self) -> Dict[str, DataKey]:
         return self._describe
 
     @AsyncStatus.wrap
     async def trigger(self) -> None:
-        """Arm the detector and wait for it to finish."""
+        # Arm the detector and wait for it to finish.
         indices_written = await self.writer.get_indices_written()
-        written_status = await self.controller.arm(DetectorTrigger.internal, num=1)
-        await written_status
-        await self.writer.wait_for_index(
-            indices_written + 1, timeout=self._frame_writing_timeout
+        written_status = await self.controller.arm(
+            num=1,
+            trigger=DetectorTrigger.internal,
         )
+        await written_status
+        end_observation = indices_written + 1
 
-    async def read(self) -> Dict[str, Reading]:
-        """Read the detector"""
-        # All data is in StreamResources, not Events, so nothing to output here
-        return {}
+        async for index in self.writer.observe_indices_written(
+            self._frame_writing_timeout
+        ):
+            if index >= end_observation:
+                break
 
-    async def collect_asset_docs(self) -> AsyncIterator[Asset]:
-        """Collect stream datum documents for all indices written."""
-        indices_written = await self.writer.get_indices_written()
+    def prepare(
+        self,
+        value: T,
+    ) -> AsyncStatus:
+        # Just arm detector for the time being
+        return AsyncStatus(self._prepare(value))
 
-        async for doc in self.writer.collect_stream_docs(indices_written):
-            yield doc
-        # async for doc in self.writer.collect_stream_docs(indices_written):
-        #     yield doc
+    async def _prepare(self, value: T) -> None:
+        """
+        Arm detector.
+
+        Prepare the detector with trigger information. This is determined at and passed
+        in from the plan level.
+
+        This currently only prepares detectors for flyscans and stepscans just use the
+        trigger information determined in trigger.
+
+        To do: Unify prepare to be use for both fly and step scans.
+
+        Args:
+            value: TriggerInfo describing how to trigger the detector
+        """
+        assert type(value) is TriggerInfo
+        self._trigger_info = value
+        self._initial_frame = await self.writer.get_indices_written()
+        self._last_frame = self._initial_frame + self._trigger_info.num
+
+        required = self.controller.get_deadtime(self._trigger_info.livetime)
+        assert required <= self._trigger_info.deadtime, (
+            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,
+        )
+        self._fly_start = time.monotonic()
 
     @AsyncStatus.wrap
-    async def unstage(self) -> None:
-        """Stop data writing."""
-        await self.writer.close()
+    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
+        ):
+            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
+
+    async def describe_collect(self) -> Dict[str, DataKey]:
+        return self._describe
+
+    async def collect_asset_docs(
+        self, index: Optional[int] = None
+    ) -> AsyncIterator[StreamAsset]:
+        # Collect stream datum documents for all indices written.
+        # The index is optional, and provided for fly scans, however this needs to be
+        # retrieved for step scans.
+        if index is None:
+            index = await self.writer.get_indices_written()
+        async for doc in self.writer.collect_stream_docs(index):
+            yield doc
+
+    async def get_index(self) -> int:
+        return await self.writer.get_indices_written()

ophyd_async/core/device.py

@@ -1,16 +1,25 @@
 """Base device"""
-from __future__ import annotations
 
 import asyncio
-import logging
 import sys
-from contextlib import suppress
-from typing import Any, Dict, Generator, Iterator, Optional, Set, Tuple, TypeVar
+from functools import cached_property
+from logging import LoggerAdapter, getLogger
+from typing import (
+    Any,
+    Coroutine,
+    Dict,
+    Generator,
+    Iterator,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+)
 
 from bluesky.protocols import HasName
 from bluesky.run_engine import call_in_bluesky_event_loop
 
-from .utils import NotConnected, wait_for_connection
+from .utils import DEFAULT_TIMEOUT, NotConnected, wait_for_connection
 
 
 class Device(HasName):
@@ -21,7 +30,10 @@ class Device(HasName):
 
     _name: str = ""
     #: The parent Device if it exists
-    parent: Optional[Device] = None
+    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)
@@ -31,7 +43,13 @@ class Device(HasName):
         """Return the name of the Device"""
         return self._name
 
-    def children(self) -> Iterator[Tuple[str, Device]]:
+    @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):
                 yield attr_name, attr
@@ -44,31 +62,68 @@ 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):
+    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) 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)
 
 
 class DeviceVector(Dict[int, VT], Device):
+    """
+    Defines device components with indices.
+
+    In the below example, foos becomes a dictionary on the parent device
+    at runtime, so parent.foos[2] returns a FooDevice. For example usage see
+    :class:`~ophyd_async.epics.demo.DynamicSensorGroup`
+    """
+
     def children(self) -> Generator[Tuple[str, Device], None, None]:
         for attr_name, attr in self.items():
             if isinstance(attr, Device):
@@ -84,9 +139,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
@@ -108,12 +163,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] = {}
@@ -140,41 +195,19 @@ class DeviceCollector:
 
     async def _on_exit(self) -> None:
         # Name and kick off connect for devices
-        tasks: Dict[asyncio.Task, str] = {}
+        connect_coroutines: Dict[str, Coroutine] = {}
         for name, obj in self._objects_on_exit.items():
             if name not in self._names_on_enter and isinstance(obj, Device):
                 if self._set_name and not obj.name:
                     obj.set_name(name)
                 if self._connect:
-                    task = asyncio.create_task(obj.connect(self._sim))
-                    tasks[task] = name
-        # Wait for all the signals to have finished
-        if tasks:
-            await self._wait_for_tasks(tasks)
-
-    async def _wait_for_tasks(self, tasks: Dict[asyncio.Task, str]):
-        done, pending = await asyncio.wait(tasks, timeout=self._timeout)
-        if pending:
-            msg = f"{len(pending)} Devices did not connect:"
-            for t in pending:
-                t.cancel()
-                with suppress(Exception):
-                    await t
-                e = t.exception()
-                msg += f"\n  {tasks[t]}: {type(e).__name__}"
-                lines = str(e).splitlines()
-                if len(lines) <= 1:
-                    msg += f": {e}"
-                else:
-                    msg += "".join(f"\n    {line}" for line in lines)
-            logging.error(msg)
-        raised = [t for t in done if t.exception()]
-        if raised:
-            logging.error(f"{len(raised)} Devices raised an error:")
-            for t in raised:
-                logging.exception(f"  {tasks[t]}:", exc_info=t.exception())
-        if pending or raised:
-            raise NotConnected("Not all Devices connected")
+                    connect_coroutines[name] = obj.connect(
+                        self._mock, timeout=self._timeout
+                    )
+
+        # Connect to all the devices
+        if connect_coroutines:
+            await wait_for_connection(**connect_coroutines)
 
     async def __aexit__(self, type, value, traceback):
         self._objects_on_exit = self._caller_locals()
@@ -182,4 +215,12 @@ class DeviceCollector:
 
     def __exit__(self, type_, value, traceback):
         self._objects_on_exit = self._caller_locals()
-        return call_in_bluesky_event_loop(self._on_exit())
+        try:
+            fut = call_in_bluesky_event_loop(self._on_exit())
+        except RuntimeError:
+            raise NotConnected(
+                "Could not connect devices. Is the bluesky event loop running? See "
+                "https://blueskyproject.io/ophyd-async/main/"
+                "user/explanations/event-loop-choice.html for more info."
+            )
+        return fut