ophyd-async 0.5.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

ophyd_async/__init__.py

@@ -1,3 +1,12 @@
+"""Top level API.
+
+.. data:: __version__
+    :type: str
+
+    Version number as calculated by https://github.com/pypa/setuptools_scm
+"""
+
+from . import core
 from ._version import __version__
 
-__all__ = ["__version__"]
+__all__ = ["__version__", "core"]

ophyd_async/__main__.py

@@ -1,16 +1,24 @@
+"""Interface for ``python -m ophyd_async``."""
+
 from argparse import ArgumentParser
+from collections.abc import Sequence
 
 from . import __version__
 
 __all__ = ["main"]
 
 
-def main(args=None):
+def main(args: Sequence[str] | None = None) -> None:
+    """Argument parser for the CLI."""
     parser = ArgumentParser()
-    parser.add_argument("-v", "--version", action="version", version=__version__)
-    args = parser.parse_args(args)
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=__version__,
+    )
+    parser.parse_args(args)
 
 
-# test with: python -m ophyd_async
 if __name__ == "__main__":
     main()

ophyd_async/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.2'
-__version_tuple__ = version_tuple = (0, 5, 2)
+__version__ = version = '0.6.0'
+__version_tuple__ = version_tuple = (0, 6, 0)

ophyd_async/core/__init__.py

@@ -61,13 +61,18 @@ from ._signal import (
     soft_signal_rw,
     wait_for_value,
 )
-from ._signal_backend import RuntimeSubsetEnum, SignalBackend, SubsetEnum
+from ._signal_backend import (
+    RuntimeSubsetEnum,
+    SignalBackend,
+    SubsetEnum,
+)
 from ._soft_signal_backend import SignalMetadata, SoftSignalBackend
 from ._status import AsyncStatus, WatchableAsyncStatus, completed_status
+from ._table import Table
 from ._utils import (
+    CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
     CalculatableTimeout,
-    CalculateTimeout,
     NotConnected,
     ReadingValueCallback,
     T,
@@ -75,6 +80,7 @@ from ._utils import (
     get_dtype,
     get_unique,
     in_micros,
+    is_pydantic_model,
     wait_for_connection,
 )
 
@@ -149,14 +155,16 @@ __all__ = [
     "WatchableAsyncStatus",
     "DEFAULT_TIMEOUT",
     "CalculatableTimeout",
-    "CalculateTimeout",
+    "CALCULATE_TIMEOUT",
     "NotConnected",
     "ReadingValueCallback",
+    "Table",
     "T",
     "WatcherUpdate",
     "get_dtype",
     "get_unique",
     "in_micros",
+    "is_pydantic_model",
     "wait_for_connection",
     "completed_status",
 ]

ophyd_async/core/_detector.py

@@ -3,21 +3,14 @@
 import asyncio
 import time
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
 from enum import Enum
 from typing import (
-    AsyncGenerator,
-    AsyncIterator,
-    Callable,
-    Dict,
     Generic,
-    List,
-    Optional,
-    Sequence,
 )
 
 from bluesky.protocols import (
     Collectable,
-    DataKey,
     Flyable,
     Preparable,
     Reading,
@@ -26,10 +19,12 @@ from bluesky.protocols import (
     Triggerable,
     WritesStreamAssets,
 )
+from event_model import DataKey
 from pydantic import BaseModel, Field
 
 from ._device import Device
 from ._protocol import AsyncConfigurable, AsyncReadable
+from ._signal import SignalR
 from ._status import AsyncStatus, WatchableAsyncStatus
 from ._utils import DEFAULT_TIMEOUT, T, WatcherUpdate, merge_gathered_dicts
 
@@ -51,20 +46,24 @@ class TriggerInfo(BaseModel):
     """Minimal set of information required to setup triggering on a detector"""
 
     #: Number of triggers that will be sent, 0 means infinite
-    number: int = Field(gt=0)
+    number: int = Field(ge=0)
     #: Sort of triggers that will be sent
-    trigger: DetectorTrigger = Field()
+    trigger: DetectorTrigger = Field(default=DetectorTrigger.internal)
     #: What is the minimum deadtime between triggers
-    deadtime: float | None = Field(ge=0)
+    deadtime: float | None = Field(default=None, ge=0)
     #: What is the maximum high time of the triggers
-    livetime: float | None = Field(ge=0)
+    livetime: float | None = Field(default=None, ge=0)
     #: What is the maximum timeout on waiting for a frame
-    frame_timeout: float | None = Field(None, gt=0)
+    frame_timeout: float | None = Field(default=None, gt=0)
     #: How many triggers make up a single StreamDatum index, to allow multiple frames
     #: from a faster detector to be zipped with a single frame from a slow detector
     #: e.g. if num=10 and multiplier=5 then the detector will take 10 frames,
     #: but publish 2 indices, and describe() will show a shape of (5, h, w)
     multiplier: int = 1
+    #: The number of times the detector can go through a complete cycle of kickoff and
+    #: complete without needing to re-arm. This is important for detectors where the
+    #: process of arming is expensive in terms of time
+    iteration: int = 1
 
 
 class DetectorControl(ABC):
@@ -78,27 +77,35 @@ class DetectorControl(ABC):
         """For a given exposure, how long should the time between exposures be"""
 
     @abstractmethod
-    async def arm(
-        self,
-        num: int,
-        trigger: DetectorTrigger = DetectorTrigger.internal,
-        exposure: Optional[float] = None,
-    ) -> AsyncStatus:
+    async def prepare(self, trigger_info: TriggerInfo):
         """
-        Arm detector, do all necessary steps to prepare detector for triggers.
+        Do all necessary steps to prepare the 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.
+            trigger_info: This is a Pydantic model which contains
+                number Expected number of frames.
+                trigger Type of trigger for which to prepare the detector. Defaults
+                to DetectorTrigger.internal.
+                livetime Livetime / 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.
+                deadtime Defaults to None. This is the minimum deadtime between
+                triggers.
+                multiplier The number of triggers grouped into a single StreamDatum
+                index.
+        """
 
-        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 arm(self) -> None:
+        """
+        Arm the detector
+        """
+
+    @abstractmethod
+    async def wait_for_idle(self):
+        """
+        This will wait on the internal _arm_status and wait for it to get disarmed/idle
         """
 
     @abstractmethod
@@ -111,7 +118,7 @@ class DetectorWriter(ABC):
     (e.g. an HDF5 file)"""
 
     @abstractmethod
-    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
+    async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
         """Open writer and wait for it to be ready for data.
 
         Args:
@@ -162,7 +169,7 @@ class StandardDetector(
         self,
         controller: DetectorControl,
         writer: DetectorWriter,
-        config_sigs: Sequence[AsyncReadable] = (),
+        config_sigs: Sequence[SignalR] = (),
         name: str = "",
     ) -> None:
         """
@@ -177,16 +184,16 @@ class StandardDetector(
         """
         self._controller = controller
         self._writer = writer
-        self._describe: Dict[str, DataKey] = {}
+        self._describe: dict[str, DataKey] = {}
         self._config_sigs = list(config_sigs)
         # For prepare
-        self._arm_status: Optional[AsyncStatus] = None
-        self._trigger_info: Optional[TriggerInfo] = None
+        self._arm_status: AsyncStatus | None = None
+        self._trigger_info: TriggerInfo | None = None
         # For kickoff
-        self._watchers: List[Callable] = []
-        self._fly_status: Optional[WatchableAsyncStatus] = None
+        self._watchers: list[Callable] = []
+        self._fly_status: WatchableAsyncStatus | None = None
         self._fly_start: float
-
+        self._iterations_completed: int = 0
         self._intial_frame: int
         self._last_frame: int
         super().__init__(name)
@@ -215,28 +222,28 @@ class StandardDetector(
                 )
             try:
                 await signal.get_value()
-            except NotImplementedError:
+            except NotImplementedError as e:
                 raise Exception(
                     f"config signal {signal.name} must be connected before it is "
                     + "passed to the detector"
-                )
+                ) from e
 
     @AsyncStatus.wrap
     async def unstage(self) -> None:
         # Stop data writing.
-        await self.writer.close()
+        await asyncio.gather(self.writer.close(), self.controller.disarm())
 
-    async def read_configuration(self) -> Dict[str, Reading]:
+    async def read_configuration(self) -> dict[str, Reading]:
         return await merge_gathered_dicts(sig.read() for sig in self._config_sigs)
 
-    async def describe_configuration(self) -> Dict[str, DataKey]:
+    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]:
+    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]:
+    async def describe(self) -> dict[str, DataKey]:
         return self._describe
 
     @AsyncStatus.wrap
@@ -248,15 +255,15 @@ class StandardDetector(
                     trigger=DetectorTrigger.internal,
                     deadtime=None,
                     livetime=None,
+                    frame_timeout=None,
                 )
             )
+        assert self._trigger_info
+        assert self._trigger_info.trigger is DetectorTrigger.internal
         # Arm the detector and wait for it to finish.
         indices_written = await self.writer.get_indices_written()
-        written_status = await self.controller.arm(
-            num=self._trigger_info.number,
-            trigger=self._trigger_info.trigger,
-        )
-        await written_status
+        await self.controller.arm()
+        await self.controller.wait_for_idle()
         end_observation = indices_written + 1
 
         async for index in self.writer.observe_indices_written(
@@ -283,35 +290,35 @@ class StandardDetector(
         Args:
             value: TriggerInfo describing how to trigger the detector
         """
-        self._trigger_info = value
         if value.trigger != DetectorTrigger.internal:
             assert (
                 value.deadtime
             ), "Deadtime must be supplied when in externally triggered mode"
         if value.deadtime:
-            required = self.controller.get_deadtime(self._trigger_info.livetime)
+            required = self.controller.get_deadtime(value.livetime)
             assert required <= value.deadtime, (
                 f"Detector {self.controller} needs at least {required}s deadtime, "
                 f"but trigger logic provides only {value.deadtime}s"
             )
+        self._trigger_info = value
         self._initial_frame = await self.writer.get_indices_written()
         self._last_frame = self._initial_frame + self._trigger_info.number
-        self._arm_status = await self.controller.arm(
-            num=self._trigger_info.number,
-            trigger=self._trigger_info.trigger,
-            exposure=self._trigger_info.livetime,
+        self._describe, _ = await asyncio.gather(
+            self.writer.open(value.multiplier), self.controller.prepare(value)
         )
-        self._fly_start = time.monotonic()
-        self._describe = await self.writer.open(value.multiplier)
+        if value.trigger != DetectorTrigger.internal:
+            await self.controller.arm()
+            self._fly_start = time.monotonic()
 
     @AsyncStatus.wrap
     async def kickoff(self):
-        if not self._arm_status:
-            raise Exception("Detector not armed!")
+        assert self._trigger_info, "Prepare must be called before kickoff!"
+        if self._iterations_completed >= self._trigger_info.iteration:
+            raise Exception(f"Kickoff called more than {self._trigger_info.iteration}")
+        self._iterations_completed += 1
 
     @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._trigger_info.frame_timeout
@@ -332,12 +339,14 @@ class StandardDetector(
             )
             if index >= self._trigger_info.number:
                 break
+        if self._iterations_completed == self._trigger_info.iteration:
+            await self.controller.wait_for_idle()
 
-    async def describe_collect(self) -> Dict[str, DataKey]:
+    async def describe_collect(self) -> dict[str, DataKey]:
         return self._describe
 
     async def collect_asset_docs(
-        self, index: Optional[int] = None
+        self, index: int | None = 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

ophyd_async/core/_device.py

@@ -2,17 +2,12 @@
 
 import asyncio
 import sys
+from collections.abc import Coroutine, Generator, Iterator
 from functools import cached_property
 from logging import LoggerAdapter, getLogger
 from typing import (
     Any,
-    Coroutine,
-    Dict,
-    Generator,
-    Iterator,
     Optional,
-    Set,
-    Tuple,
     TypeVar,
 )
 
@@ -32,7 +27,7 @@ class Device(HasName):
     #: 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_task: asyncio.Task | None = None
 
     # Used to check if the previous connect was mocked,
     # if the next mock value differs then we fail
@@ -52,7 +47,7 @@ class Device(HasName):
             getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
         )
 
-    def children(self) -> Iterator[Tuple[str, "Device"]]:
+    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
@@ -127,7 +122,7 @@ class Device(HasName):
 VT = TypeVar("VT", bound=Device)
 
 
-class DeviceVector(Dict[int, VT], Device):
+class DeviceVector(dict[int, VT], Device):
     """
     Defines device components with indices.
 
@@ -136,7 +131,7 @@ class DeviceVector(Dict[int, VT], Device):
     :class:`~ophyd_async.epics.demo.DynamicSensorGroup`
     """
 
-    def children(self) -> Generator[Tuple[str, Device], None, None]:
+    def children(self) -> Generator[tuple[str, Device], None, None]:
         for attr_name, attr in self.items():
             if isinstance(attr, Device):
                 yield str(attr_name), attr
@@ -182,8 +177,8 @@ class DeviceCollector:
         self._connect = connect
         self._mock = mock
         self._timeout = timeout
-        self._names_on_enter: Set[str] = set()
-        self._objects_on_exit: Dict[str, Any] = {}
+        self._names_on_enter: set[str] = set()
+        self._objects_on_exit: dict[str, Any] = {}
 
     def _caller_locals(self):
         """Walk up until we find a stack frame that doesn't have us as self"""
@@ -195,6 +190,9 @@ class DeviceCollector:
             caller_frame = tb.tb_frame
             while caller_frame.f_locals.get("self", None) is self:
                 caller_frame = caller_frame.f_back
+                assert (
+                    caller_frame
+                ), "No previous frame to the one with self in it, this shouldn't happen"
             return caller_frame.f_locals
 
     def __enter__(self) -> "DeviceCollector":
@@ -207,7 +205,7 @@ class DeviceCollector:
 
     async def _on_exit(self) -> None:
         # Name and kick off connect for devices
-        connect_coroutines: Dict[str, Coroutine] = {}
+        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:
@@ -229,10 +227,10 @@ class DeviceCollector:
         self._objects_on_exit = self._caller_locals()
         try:
             fut = call_in_bluesky_event_loop(self._on_exit())
-        except RuntimeError:
+        except RuntimeError as e:
             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."
-            )
+            ) from e
         return fut

ophyd_async/core/_device_save_loader.py

@@ -1,5 +1,7 @@
+from collections.abc import Callable, Generator, Sequence
 from enum import Enum
-from typing import Any, Callable, Dict, Generator, List, Optional, Sequence
+from pathlib import Path
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
@@ -7,6 +9,7 @@ import yaml
 from bluesky.plan_stubs import abs_set, wait
 from bluesky.protocols import Location
 from bluesky.utils import Msg
+from pydantic import BaseModel
 
 from ._device import Device
 from ._signal import SignalRW
@@ -18,16 +21,22 @@ def ndarray_representer(dumper: yaml.Dumper, array: npt.NDArray[Any]) -> yaml.No
     )
 
 
+def pydantic_model_abstraction_representer(
+    dumper: yaml.Dumper, model: BaseModel
+) -> yaml.Node:
+    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(OphydDumper, self).represent_data(data)
+        return super().represent_data(data)
 
 
 def get_signal_values(
-    signals: Dict[str, SignalRW[Any]], ignore: Optional[List[str]] = None
-) -> Generator[Msg, Sequence[Location[Any]], Dict[str, Any]]:
+    signals: dict[str, SignalRW[Any]], ignore: list[str] | None = None
+) -> Generator[Msg, Sequence[Location[Any]], dict[str, Any]]:
     """Get signal values in bulk.
 
     Used as part of saving the signals of a device to a yaml file.
@@ -59,13 +68,10 @@ def get_signal_values(
     }
     selected_values = yield Msg("locate", *selected_signals.values())
 
-    # TODO: investigate wrong type hints
-    if isinstance(selected_values, dict):
-        selected_values = [selected_values]  # type: ignore
-
     assert selected_values is not None, "No signalRW's were able to be located"
     named_values = {
-        key: value["setpoint"] for key, value in zip(selected_signals, selected_values)
+        key: value["setpoint"]
+        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})
@@ -73,8 +79,8 @@ def get_signal_values(
 
 
 def walk_rw_signals(
-    device: Device, path_prefix: Optional[str] = ""
-) -> Dict[str, SignalRW[Any]]:
+    device: Device, path_prefix: str | None = ""
+) -> dict[str, SignalRW[Any]]:
     """Retrieve all SignalRWs from a device.
 
     Stores retrieved signals with their dotted attribute paths in a dictionary. Used as
@@ -104,7 +110,7 @@ def walk_rw_signals(
     if not path_prefix:
         path_prefix = ""
 
-    signals: Dict[str, SignalRW[Any]] = {}
+    signals: dict[str, SignalRW[Any]] = {}
     for attr_name, attr in device.children():
         dot_path = f"{path_prefix}{attr_name}"
         if type(attr) is SignalRW:
@@ -114,7 +120,7 @@ def walk_rw_signals(
     return signals
 
 
-def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
+def save_to_yaml(phases: Sequence[dict[str, Any]], save_path: str | Path) -> None:
     """Plan which serialises a phase or set of phases of SignalRWs to a yaml file.
 
     Parameters
@@ -134,12 +140,17 @@ def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
     """
 
     yaml.add_representer(np.ndarray, ndarray_representer, Dumper=yaml.Dumper)
+    yaml.add_multi_representer(
+        BaseModel,
+        pydantic_model_abstraction_representer,
+        Dumper=yaml.Dumper,
+    )
 
     with open(save_path, "w") as file:
         yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)
 
 
-def load_from_yaml(save_path: str) -> Sequence[Dict[str, Any]]:
+def load_from_yaml(save_path: str) -> Sequence[dict[str, Any]]:
     """Plan that returns a list of dicts with saved signal values from a yaml file.
 
     Parameters
@@ -152,12 +163,12 @@ def load_from_yaml(save_path: str) -> Sequence[Dict[str, Any]]:
     :func:`ophyd_async.core.save_to_yaml`
     :func:`ophyd_async.core.set_signal_values`
     """
-    with open(save_path, "r") as file:
+    with open(save_path) as file:
         return yaml.full_load(file)
 
 
 def set_signal_values(
-    signals: Dict[str, SignalRW[Any]], values: Sequence[Dict[str, Any]]
+    signals: dict[str, SignalRW[Any]], values: Sequence[dict[str, Any]]
 ) -> Generator[Msg, None, None]:
     """Maps signals from a yaml file into device signals.
 
@@ -217,7 +228,7 @@ def load_device(device: Device, path: str):
     yield from set_signal_values(signals_to_set, values)
 
 
-def all_at_once(values: Dict[str, Any]) -> Sequence[Dict[str, Any]]:
+def all_at_once(values: dict[str, Any]) -> Sequence[dict[str, Any]]:
     """Sort all the values into a single phase so they are set all at once"""
     return [values]
 
@@ -225,8 +236,8 @@ def all_at_once(values: Dict[str, Any]) -> Sequence[Dict[str, Any]]:
 def save_device(
     device: Device,
     path: str,
-    sorter: Callable[[Dict[str, Any]], Sequence[Dict[str, Any]]] = all_at_once,
-    ignore: Optional[List[str]] = None,
+    sorter: Callable[[dict[str, Any]], Sequence[dict[str, Any]]] = all_at_once,
+    ignore: list[str] | None = None,
 ):
     """Plan that saves the state of all PV's on a device using a sorter.
 

ophyd_async/core/_flyer.py

@@ -1,7 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Generic, Sequence
+from collections.abc import Sequence
+from typing import Generic
 
-from bluesky.protocols import DataKey, Flyable, Preparable, Reading, Stageable
+from bluesky.protocols import Flyable, Preparable, Reading, Stageable
+from event_model import DataKey
 
 from ._device import Device
 from ._signal import SignalR
@@ -72,12 +74,12 @@ class StandardFlyer(
     async def complete(self) -> None:
         await self._trigger_logic.complete()
 
-    async def describe_configuration(self) -> Dict[str, DataKey]:
+    async def describe_configuration(self) -> dict[str, DataKey]:
         return await merge_gathered_dicts(
             [sig.describe() for sig in self._configuration_signals]
         )
 
-    async def read_configuration(self) -> Dict[str, Reading]:
+    async def read_configuration(self) -> dict[str, Reading]:
         return await merge_gathered_dicts(
             [sig.read() for sig in self._configuration_signals]
         )