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

ophyd_async/core/_table.py

@@ -0,0 +1,146 @@
+from enum import Enum
+from typing import TypeVar, get_args, get_origin
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, model_validator
+
+TableSubclass = TypeVar("TableSubclass", bound="Table")
+
+
+def _concat(value1, value2):
+    if isinstance(value1, np.ndarray):
+        return np.concatenate((value1, value2))
+    else:
+        return value1 + value2
+
+
+class Table(BaseModel):
+    """An abstraction of a Table of str to numpy array."""
+
+    model_config = ConfigDict(validate_assignment=True, strict=False)
+
+    @staticmethod
+    def row(cls: type[TableSubclass], **kwargs) -> TableSubclass:  # type: ignore
+        arrayified_kwargs = {}
+        for field_name, field_value in cls.model_fields.items():
+            value = kwargs.pop(field_name)
+            if field_value.default_factory is None:
+                raise ValueError(
+                    "`Table` models should have default factories for their "
+                    "mutable empty columns."
+                )
+            default_array = field_value.default_factory()
+            if isinstance(default_array, np.ndarray):
+                arrayified_kwargs[field_name] = np.array(
+                    [value], dtype=default_array.dtype
+                )
+            elif issubclass(type(value), Enum) and isinstance(value, str):
+                arrayified_kwargs[field_name] = [value]
+            else:
+                raise TypeError(
+                    "Row column should be numpy arrays or sequence of string `Enum`."
+                )
+        if kwargs:
+            raise TypeError(
+                f"Unexpected keyword arguments {kwargs.keys()} for {cls.__name__}."
+            )
+        return cls(**arrayified_kwargs)
+
+    def __add__(self, right: TableSubclass) -> TableSubclass:
+        """Concatenate the arrays in field values."""
+
+        if type(right) is not type(self):
+            raise RuntimeError(
+                f"{right} is not a `Table`, or is not the same "
+                f"type of `Table` as {self}."
+            )
+
+        return type(right)(
+            **{
+                field_name: _concat(
+                    getattr(self, field_name), getattr(right, field_name)
+                )
+                for field_name in self.model_fields
+            }
+        )
+
+    def numpy_dtype(self) -> np.dtype:
+        dtype = []
+        for field_name, field_value in self.model_fields.items():
+            if np.ndarray in (
+                get_origin(field_value.annotation),
+                field_value.annotation,
+            ):
+                dtype.append((field_name, getattr(self, field_name).dtype))
+            else:
+                enum_type = get_args(field_value.annotation)[0]
+                assert issubclass(enum_type, Enum)
+                enum_values = [element.value for element in enum_type]
+                max_length_in_enum = max(len(value) for value in enum_values)
+                dtype.append((field_name, np.dtype(f"<U{max_length_in_enum}")))
+
+        return np.dtype(dtype)
+
+    def numpy_table(self):
+        # It would be nice to be able to use np.transpose for this,
+        # but it defaults to the largest dtype for everything.
+        dtype = self.numpy_dtype()
+        transposed_list = [
+            np.array(tuple(row), dtype=dtype)
+            for row in zip(*self.numpy_columns(), strict=False)
+        ]
+        transposed = np.array(transposed_list, dtype=dtype)
+        return transposed
+
+    def numpy_columns(self) -> list[np.ndarray]:
+        """Columns in the table can be lists of string enums or numpy arrays.
+
+        This method returns the columns, converting the string enums to numpy arrays.
+        """
+
+        columns = []
+        for field_name, field_value in self.model_fields.items():
+            if np.ndarray in (
+                get_origin(field_value.annotation),
+                field_value.annotation,
+            ):
+                columns.append(getattr(self, field_name))
+            else:
+                enum_type = get_args(field_value.annotation)[0]
+                assert issubclass(enum_type, Enum)
+                enum_values = [element.value for element in enum_type]
+                max_length_in_enum = max(len(value) for value in enum_values)
+                dtype = np.dtype(f"<U{max_length_in_enum}")
+
+                columns.append(
+                    np.array(
+                        [enum.value for enum in getattr(self, field_name)], dtype=dtype
+                    )
+                )
+
+        return columns
+
+    @model_validator(mode="after")
+    def validate_arrays(self) -> "Table":
+        first_length = len(next(iter(self))[1])
+        assert all(
+            len(field_value) == first_length for _, field_value in self
+        ), "Rows should all be of equal size."
+
+        if not all(
+            # Checks if the values are numpy subtypes if the array is a numpy array,
+            # or if the value is a string enum.
+            np.issubdtype(getattr(self, field_name).dtype, default_array.dtype)
+            if isinstance(
+                default_array := self.model_fields[field_name].default_factory(),  # type: ignore
+                np.ndarray,
+            )
+            else issubclass(get_args(field_value.annotation)[0], Enum)
+            for field_name, field_value in self.model_fields.items()
+        ):
+            raise ValueError(
+                f"Cannot construct a `{type(self).__name__}`, "
+                "some rows have incorrect types."
+            )
+
+        return self

ophyd_async/core/_utils.py

@@ -2,23 +2,13 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from collections.abc import Awaitable, Callable, Iterable
 from dataclasses import dataclass
-from typing import (
-    Awaitable,
-    Callable,
-    Dict,
-    Generic,
-    Iterable,
-    List,
-    Optional,
-    ParamSpec,
-    Type,
-    TypeVar,
-    Union,
-)
+from typing import Generic, Literal, ParamSpec, TypeVar, get_origin
 
 import numpy as np
 from bluesky.protocols import Reading
+from pydantic import BaseModel
 
 T = TypeVar("T")
 P = ParamSpec("P")
@@ -28,18 +18,18 @@ Callback = Callable[[T], None]
 #: monitor updates
 ReadingValueCallback = Callable[[Reading, T], None]
 DEFAULT_TIMEOUT = 10.0
-ErrorText = Union[str, Dict[str, Exception]]
+ErrorText = str | dict[str, Exception]
 
 
-class CalculateTimeout:
-    """Sentinel class used to implement ``myfunc(timeout=CalculateTimeout)``
+CALCULATE_TIMEOUT = "CALCULATE_TIMEOUT"
+"""Sentinel used to implement ``myfunc(timeout=CalculateTimeout)``
 
-    This signifies that the function should calculate a suitable non-zero
-    timeout itself
-    """
+This signifies that the function should calculate a suitable non-zero
+timeout itself
+"""
 
 
-CalculatableTimeout = float | None | Type[CalculateTimeout]
+CalculatableTimeout = float | None | Literal["CALCULATE_TIMEOUT"]
 
 
 class NotConnected(Exception):
@@ -115,7 +105,7 @@ async def wait_for_connection(**coros: Awaitable[None]):
     results = await asyncio.gather(*coros.values(), return_exceptions=True)
     exceptions = {}
 
-    for name, result in zip(coros, results):
+    for name, result in zip(coros, results, strict=False):
         if isinstance(result, Exception):
             exceptions[name] = result
             if not isinstance(result, NotConnected):
@@ -129,7 +119,7 @@ async def wait_for_connection(**coros: Awaitable[None]):
         raise NotConnected(exceptions)
 
 
-def get_dtype(typ: Type) -> Optional[np.dtype]:
+def get_dtype(typ: type) -> np.dtype | None:
     """Get the runtime dtype from a numpy ndarray type annotation
 
     >>> import numpy.typing as npt
@@ -144,8 +134,8 @@ def get_dtype(typ: Type) -> Optional[np.dtype]:
     return None
 
 
-def get_unique(values: Dict[str, T], types: str) -> T:
-    """If all values are the same, return that value, otherwise return TypeError
+def get_unique(values: dict[str, T], types: str) -> T:
+    """If all values are the same, return that value, otherwise raise TypeError
 
     >>> get_unique({"a": 1, "b": 1}, "integers")
     1
@@ -162,21 +152,21 @@ def get_unique(values: Dict[str, T], types: str) -> T:
 
 
 async def merge_gathered_dicts(
-    coros: Iterable[Awaitable[Dict[str, T]]],
-) -> Dict[str, T]:
+    coros: Iterable[Awaitable[dict[str, T]]],
+) -> dict[str, T]:
     """Merge dictionaries produced by a sequence of coroutines.
 
     Can be used for merging ``read()`` or ``describe``. For instance::
 
         combined_read = await merge_gathered_dicts(s.read() for s in signals)
     """
-    ret: Dict[str, T] = {}
+    ret: dict[str, T] = {}
     for result in await asyncio.gather(*coros):
         ret.update(result)
     return ret
 
 
-async def gather_list(coros: Iterable[Awaitable[T]]) -> List[T]:
+async def gather_list(coros: Iterable[Awaitable[T]]) -> list[T]:
     return await asyncio.gather(*coros)
 
 
@@ -195,3 +185,9 @@ def in_micros(t: float) -> int:
     if t < 0:
         raise ValueError(f"Expected a positive time in seconds, got {t!r}")
     return int(np.ceil(t * 1e6))
+
+
+def is_pydantic_model(datatype) -> bool:
+    while origin := get_origin(datatype):
+        datatype = origin
+    return datatype and issubclass(datatype, BaseModel)

ophyd_async/epics/adaravis/_aravis_controller.py

@@ -1,12 +1,13 @@
 import asyncio
-from typing import Literal, Optional, Tuple
+from typing import Literal
 
 from ophyd_async.core import (
-    AsyncStatus,
-    DetectorControl,
+    DetectorController,
     DetectorTrigger,
+    TriggerInfo,
     set_and_wait_for_value,
 )
+from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 from ._aravis_io import AravisDriverIO, AravisTriggerMode, AravisTriggerSource
@@ -17,45 +18,45 @@ from ._aravis_io import AravisDriverIO, AravisTriggerMode, AravisTriggerSource
 _HIGHEST_POSSIBLE_DEADTIME = 1961e-6
 
 
-class AravisController(DetectorControl):
+class AravisController(DetectorController):
     GPIO_NUMBER = Literal[1, 2, 3, 4]
 
     def __init__(self, driver: AravisDriverIO, gpio_number: GPIO_NUMBER) -> None:
         self._drv = driver
         self.gpio_number = gpio_number
+        self._arm_status: AsyncStatus | None = None
 
-    def get_deadtime(self, exposure: float) -> float:
+    def get_deadtime(self, exposure: float | None) -> float:
         return _HIGHEST_POSSIBLE_DEADTIME
 
-    async def arm(
-        self,
-        num: int = 0,
-        trigger: DetectorTrigger = DetectorTrigger.internal,
-        exposure: Optional[float] = None,
-    ) -> AsyncStatus:
-        if num == 0:
+    async def prepare(self, trigger_info: TriggerInfo):
+        if trigger_info.total_number_of_triggers == 0:
             image_mode = adcore.ImageMode.continuous
         else:
             image_mode = adcore.ImageMode.multiple
-        if exposure is not None:
+        if (exposure := trigger_info.livetime) is not None:
             await self._drv.acquire_time.set(exposure)
 
-        trigger_mode, trigger_source = self._get_trigger_info(trigger)
+        trigger_mode, trigger_source = self._get_trigger_info(trigger_info.trigger)
         # trigger mode must be set first and on it's own!
         await self._drv.trigger_mode.set(trigger_mode)
 
         await asyncio.gather(
             self._drv.trigger_source.set(trigger_source),
-            self._drv.num_images.set(num),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
             self._drv.image_mode.set(image_mode),
         )
 
-        status = await set_and_wait_for_value(self._drv.acquire, True)
-        return status
+    async def arm(self):
+        self._arm_status = await set_and_wait_for_value(self._drv.acquire, True)
+
+    async def wait_for_idle(self):
+        if self._arm_status:
+            await self._arm_status
 
     def _get_trigger_info(
         self, trigger: DetectorTrigger
-    ) -> Tuple[AravisTriggerMode, AravisTriggerSource]:
+    ) -> tuple[AravisTriggerMode, AravisTriggerSource]:
         supported_trigger_types = (
             DetectorTrigger.constant_gate,
             DetectorTrigger.edge_trigger,
@@ -70,7 +71,7 @@ class AravisController(DetectorControl):
         if trigger == DetectorTrigger.internal:
             return AravisTriggerMode.off, "Freerun"
         else:
-            return (AravisTriggerMode.on, f"Line{self.gpio_number}")
+            return (AravisTriggerMode.on, f"Line{self.gpio_number}")  # type: ignore
 
     async def disarm(self):
         await adcore.stop_busy_record(self._drv.acquire, False, timeout=1)

ophyd_async/epics/adaravis/_aravis_io.py

@@ -38,6 +38,7 @@ class AravisDriverIO(adcore.ADBaseIO):
             AravisTriggerMode, prefix + "TriggerMode"
         )
         self.trigger_source = epics_signal_rw_rbv(
-            AravisTriggerSource, prefix + "TriggerSource"
+            AravisTriggerSource,  # type: ignore
+            prefix + "TriggerSource",
         )
         super().__init__(prefix, name=name)

ophyd_async/epics/adcore/_core_io.py

@@ -135,4 +135,6 @@ class NDFileHDFIO(NDPluginBaseIO):
         self.array_size0 = epics_signal_r(int, prefix + "ArraySize0")
         self.array_size1 = epics_signal_r(int, prefix + "ArraySize1")
         self.create_directory = epics_signal_rw(int, prefix + "CreateDirectory")
+        self.num_frames_chunks = epics_signal_r(int, prefix + "NumFramesChunks_RBV")
+        self.chunk_size_auto = epics_signal_rw_rbv(bool, prefix + "ChunkSizeAuto")
         super().__init__(prefix, name)

ophyd_async/epics/adcore/_core_logic.py

@@ -1,11 +1,10 @@
 import asyncio
-from typing import FrozenSet, Set
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     AsyncStatus,
     DatasetDescriber,
-    DetectorControl,
+    DetectorController,
     set_and_wait_for_value,
 )
 from ophyd_async.epics.adcore._utils import convert_ad_dtype_to_np
@@ -14,7 +13,7 @@ from ._core_io import ADBaseIO, DetectorState
 
 # Default set of states that we should consider "good" i.e. the acquisition
 #  is complete and went well
-DEFAULT_GOOD_STATES: FrozenSet[DetectorState] = frozenset(
+DEFAULT_GOOD_STATES: frozenset[DetectorState] = frozenset(
     [DetectorState.Idle, DetectorState.Aborted]
 )
 
@@ -35,7 +34,7 @@ class ADBaseDatasetDescriber(DatasetDescriber):
 
 
 async def set_exposure_time_and_acquire_period_if_supplied(
-    controller: DetectorControl,
+    controller: DetectorController,
     driver: ADBaseIO,
     exposure: float | None = None,
     timeout: float = DEFAULT_TIMEOUT,
@@ -66,7 +65,7 @@ async def set_exposure_time_and_acquire_period_if_supplied(
 
 async def start_acquiring_driver_and_ensure_status(
     driver: ADBaseIO,
-    good_states: Set[DetectorState] = set(DEFAULT_GOOD_STATES),
+    good_states: frozenset[DetectorState] = frozenset(DEFAULT_GOOD_STATES),
     timeout: float = DEFAULT_TIMEOUT,
 ) -> AsyncStatus:
     """

ophyd_async/epics/adcore/_hdf_writer.py

@@ -1,9 +1,10 @@
 import asyncio
+from collections.abc import AsyncGenerator, AsyncIterator
 from pathlib import Path
-from typing import AsyncGenerator, AsyncIterator, Dict, List, Optional
 from xml.etree import ElementTree as ET
 
-from bluesky.protocols import DataKey, Hints, StreamAsset
+from bluesky.protocols import Hints, StreamAsset
+from event_model import DataKey
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
@@ -42,19 +43,22 @@ class ADHDFWriter(DetectorWriter):
         self._dataset_describer = dataset_describer
 
         self._plugins = plugins
-        self._capture_status: Optional[AsyncStatus] = None
-        self._datasets: List[HDFDataset] = []
-        self._file: Optional[HDFFile] = None
+        self._capture_status: AsyncStatus | None = None
+        self._datasets: list[HDFDataset] = []
+        self._file: HDFFile | None = None
         self._multiplier = 1
 
-    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
+    async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
         self._file = None
-        info = self._path_provider(device_name=self.hdf.name)
+        info = self._path_provider(device_name=self._name_provider())
 
         # Set the directory creation depth first, since dir creation callback happens
         # when directory path PV is processed.
         await self.hdf.create_directory.set(info.create_dir_depth)
 
+        # Make sure we are using chunk auto-sizing
+        await asyncio.gather(self.hdf.chunk_size_auto.set(True))
+
         await asyncio.gather(
             self.hdf.num_extra_dims.set(0),
             self.hdf.lazy_open.set(True),
@@ -83,6 +87,9 @@ class ADHDFWriter(DetectorWriter):
         self._multiplier = multiplier
         outer_shape = (multiplier,) if multiplier > 1 else ()
 
+        # Determine number of frames that will be saved per HDF chunk
+        frames_per_chunk = await self.hdf.num_frames_chunks.get_value()
+
         # Add the main data
         self._datasets = [
             HDFDataset(
@@ -91,6 +98,7 @@ class ADHDFWriter(DetectorWriter):
                 shape=detector_shape,
                 dtype_numpy=np_dtype,
                 multiplier=multiplier,
+                chunk_shape=(frames_per_chunk, *detector_shape),
             )
         ]
         # And all the scalar datasets
@@ -117,6 +125,9 @@ class ADHDFWriter(DetectorWriter):
                             (),
                             np_datatype,
                             multiplier,
+                            # NDAttributes appear to always be configured with
+                            # this chunk size
+                            chunk_shape=(16384,),
                         )
                     )
 
@@ -125,7 +136,7 @@ class ADHDFWriter(DetectorWriter):
                 source=self.hdf.full_file_name.source,
                 shape=outer_shape + tuple(ds.shape),
                 dtype="array" if ds.shape else "number",
-                dtype_numpy=ds.dtype_numpy,
+                dtype_numpy=ds.dtype_numpy,  # type: ignore
                 external="STREAM:",
             )
             for ds in self._datasets

ophyd_async/epics/adcore/_single_trigger.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Sequence
+from collections.abc import Sequence
 
 from bluesky.protocols import Triggerable
 

ophyd_async/epics/adcore/_utils.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional
 
 from ophyd_async.core import DEFAULT_TIMEOUT, SignalRW, T, wait_for_value
 from ophyd_async.core._signal import SignalR
@@ -51,8 +50,8 @@ def convert_pv_dtype_to_np(datatype: str) -> str:
     else:
         try:
             np_datatype = convert_ad_dtype_to_np(_pvattribute_to_ad_datatype[datatype])
-        except KeyError:
-            raise ValueError(f"Invalid dbr type {datatype}")
+        except KeyError as e:
+            raise ValueError(f"Invalid dbr type {datatype}") from e
     return np_datatype
 
 
@@ -69,8 +68,8 @@ def convert_param_dtype_to_np(datatype: str) -> str:
             np_datatype = convert_ad_dtype_to_np(
                 _paramattribute_to_ad_datatype[datatype]
             )
-        except KeyError:
-            raise ValueError(f"Invalid datatype {datatype}")
+        except KeyError as e:
+            raise ValueError(f"Invalid datatype {datatype}") from e
     return np_datatype
 
 
@@ -126,7 +125,7 @@ async def stop_busy_record(
     signal: SignalRW[T],
     value: T,
     timeout: float = DEFAULT_TIMEOUT,
-    status_timeout: Optional[float] = None,
+    status_timeout: float | None = None,
 ) -> None:
     await signal.set(value, wait=False, timeout=status_timeout)
     await wait_for_value(signal, value, timeout=timeout)

ophyd_async/epics/adkinetix/_kinetix_controller.py

@@ -1,7 +1,8 @@
 import asyncio
-from typing import Optional
 
-from ophyd_async.core import AsyncStatus, DetectorControl, DetectorTrigger
+from ophyd_async.core import DetectorController, DetectorTrigger
+from ophyd_async.core._detector import TriggerInfo
+from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 from ._kinetix_io import KinetixDriverIO, KinetixTriggerMode
@@ -14,33 +15,37 @@ KINETIX_TRIGGER_MODE_MAP = {
 }
 
 
-class KinetixController(DetectorControl):
+class KinetixController(DetectorController):
     def __init__(
         self,
         driver: KinetixDriverIO,
     ) -> None:
         self._drv = driver
+        self._arm_status: AsyncStatus | None = None
 
-    def get_deadtime(self, exposure: float) -> float:
+    def get_deadtime(self, exposure: float | None) -> float:
         return 0.001
 
-    async def arm(
-        self,
-        num: int,
-        trigger: DetectorTrigger = DetectorTrigger.internal,
-        exposure: Optional[float] = None,
-    ) -> AsyncStatus:
+    async def prepare(self, trigger_info: TriggerInfo):
         await asyncio.gather(
-            self._drv.trigger_mode.set(KINETIX_TRIGGER_MODE_MAP[trigger]),
-            self._drv.num_images.set(num),
+            self._drv.trigger_mode.set(KINETIX_TRIGGER_MODE_MAP[trigger_info.trigger]),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
             self._drv.image_mode.set(adcore.ImageMode.multiple),
         )
-        if exposure is not None and trigger not in [
+        if trigger_info.livetime is not None and trigger_info.trigger not in [
             DetectorTrigger.variable_gate,
             DetectorTrigger.constant_gate,
         ]:
-            await self._drv.acquire_time.set(exposure)
-        return await adcore.start_acquiring_driver_and_ensure_status(self._drv)
+            await self._drv.acquire_time.set(trigger_info.livetime)
+
+    async def arm(self):
+        self._arm_status = await adcore.start_acquiring_driver_and_ensure_status(
+            self._drv
+        )
+
+    async def wait_for_idle(self):
+        if self._arm_status:
+            await self._arm_status
 
     async def disarm(self):
         await adcore.stop_busy_record(self._drv.acquire, False, timeout=1)

ophyd_async/epics/adpilatus/_pilatus_controller.py

@@ -1,19 +1,19 @@
 import asyncio
-from typing import Optional
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
-    AsyncStatus,
-    DetectorControl,
+    DetectorController,
     DetectorTrigger,
     wait_for_value,
 )
+from ophyd_async.core._detector import TriggerInfo
+from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 from ._pilatus_io import PilatusDriverIO, PilatusTriggerMode
 
 
-class PilatusController(DetectorControl):
+class PilatusController(DetectorController):
     _supported_trigger_types = {
         DetectorTrigger.internal: PilatusTriggerMode.internal,
         DetectorTrigger.constant_gate: PilatusTriggerMode.ext_enable,
@@ -27,29 +27,31 @@ class PilatusController(DetectorControl):
     ) -> None:
         self._drv = driver
         self._readout_time = readout_time
+        self._arm_status: AsyncStatus | None = None
 
-    def get_deadtime(self, exposure: float) -> float:
+    def get_deadtime(self, exposure: float | None) -> float:
         return self._readout_time
 
-    async def arm(
-        self,
-        num: int,
-        trigger: DetectorTrigger = DetectorTrigger.internal,
-        exposure: Optional[float] = None,
-    ) -> AsyncStatus:
-        if exposure is not None:
+    async def prepare(self, trigger_info: TriggerInfo):
+        if trigger_info.livetime is not None:
             await adcore.set_exposure_time_and_acquire_period_if_supplied(
-                self, self._drv, exposure
+                self, self._drv, trigger_info.livetime
             )
         await asyncio.gather(
-            self._drv.trigger_mode.set(self._get_trigger_mode(trigger)),
-            self._drv.num_images.set(999_999 if num == 0 else num),
+            self._drv.trigger_mode.set(self._get_trigger_mode(trigger_info.trigger)),
+            self._drv.num_images.set(
+                999_999
+                if trigger_info.total_number_of_triggers == 0
+                else trigger_info.total_number_of_triggers
+            ),
             self._drv.image_mode.set(adcore.ImageMode.multiple),
         )
 
+    async def arm(self):
         # Standard arm the detector and wait for the acquire PV to be True
-        idle_status = await adcore.start_acquiring_driver_and_ensure_status(self._drv)
-
+        self._arm_status = await adcore.start_acquiring_driver_and_ensure_status(
+            self._drv
+        )
         # The pilatus has an additional PV that goes True when the camserver
         # is actually ready. Should wait for that too or we risk dropping
         # a frame
@@ -59,7 +61,9 @@ class PilatusController(DetectorControl):
             timeout=DEFAULT_TIMEOUT,
         )
 
-        return idle_status
+    async def wait_for_idle(self):
+        if self._arm_status:
+            await self._arm_status
 
     @classmethod
     def _get_trigger_mode(cls, trigger: DetectorTrigger) -> PilatusTriggerMode: