Merge branch 'main' into add_panda_save_load

.github/workflows/code.yml

@@ -111,6 +111,7 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       id-token: write
+      contents: write
     env:
       HAS_PYPI_TOKEN: ${{ secrets.PYPI_TOKEN != '' }}
 

README.rst

@@ -1,7 +1,7 @@
 Ophyd Async
 ===========
 
-|build_status| |coverage| |pypi_version| |license|
+|code_ci| |docs_ci| |coverage| |pypi_version| |license|
 
 Asynchronous device abstraction framework, building on `Ophyd`_.
 
@@ -35,16 +35,20 @@ NOTE: ophyd-async is included on a provisional basis until the v1.0 release.
 
 See the tutorials for usage examples.
 
-.. |build_status| image:: https://github.com/bluesky/ophyd/workflows/Unit%20Tests/badge.svg?branch=master
-    :target: https://github.com/bluesky/ophyd/actions?query=workflow%3A%22Unit+Tests%22
-    :alt: Build Status
+.. |code_ci| image:: https://github.com/bluesky/ophyd-async/actions/workflows/code.yml/badge.svg?branch=main
+    :target: https://github.com/bluesky/ophyd-async/actions/workflows/code.yml
+    :alt: Code CI
 
-.. |coverage| image:: https://codecov.io/gh/bluesky/ophyd/branch/master/graph/badge.svg
-    :target: https://codecov.io/gh/bluesky/ophyd
+.. |docs_ci| image:: https://github.com/bluesky/ophyd-async/actions/workflows/docs.yml/badge.svg?branch=main
+    :target: https://github.com/bluesky/ophyd-async/actions/workflows/docs.yml
+    :alt: Docs CI
+
+.. |coverage| image:: https://codecov.io/gh/bluesky/ophyd-async/branch/master/graph/badge.svg
+    :target: https://codecov.io/gh/bluesky/ophyd-async
     :alt: Test Coverage
 
-.. |pypi_version| image:: https://img.shields.io/pypi/v/ophyd.svg
-    :target: https://pypi.org/project/ophyd
+.. |pypi_version| image:: https://img.shields.io/pypi/v/ophyd-async.svg
+    :target: https://pypi.org/project/ophyd-async
     :alt: Latest PyPI version
 
 .. |license| image:: https://img.shields.io/badge/License-BSD%203--Clause-blue.svg

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "numpy",
     "packaging",
     "pint",
-    "bluesky @ git+https://github.com/bluesky/bluesky.git",
+    "bluesky",
     "event-model",
     "p4p",
     "pyyaml",

src/ophyd_async/core/device_save_loader.py

@@ -1,23 +1,40 @@
 from enum import Enum
-from typing import Any, Callable, Dict, Generator, List, Optional, Sequence
+from functools import partial
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
 import yaml
 from bluesky.plan_stubs import abs_set, wait
 from bluesky.protocols import Location
 from bluesky.utils import Msg
+from epicscorelibs.ca.dbr import ca_array, ca_float, ca_int, ca_str
 
 from .device import Device
 from .signal import SignalRW
 
+CaType = Union[ca_float, ca_int, ca_str, ca_array]
+
 
 def ndarray_representer(dumper: yaml.Dumper, array: npt.NDArray[Any]) -> yaml.Node:
     return dumper.represent_sequence(
         "tag:yaml.org,2002:seq", array.tolist(), flow_style=True
     )
 
 
+def ca_dbr_representer(dumper: yaml.Dumper, value: CaType) -> yaml.Node:
+    # if it's an array, just call ndarray_representer...
+    represent_array = partial(ndarray_representer, dumper)
+
+    representers: Dict[CaType, Callable[[CaType], yaml.Node]] = {
+        ca_float: dumper.represent_float,
+        ca_int: dumper.represent_int,
+        ca_str: dumper.represent_str,
+        ca_array: represent_array,
+    }
+    return representers[type(value)](value)
+
+
 class OphydDumper(yaml.Dumper):
     def represent_data(self, data: Any) -> Any:
         if isinstance(data, Enum):
@@ -58,6 +75,11 @@ def get_signal_values(
         key: signal for key, signal in signals.items() if key not in ignore
     }
     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)
@@ -127,7 +149,14 @@ def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
     :func:`ophyd_async.core.get_signal_values`
     :func:`ophyd_async.core.load_from_yaml`
     """
+
     yaml.add_representer(np.ndarray, ndarray_representer, Dumper=yaml.Dumper)
+
+    yaml.add_representer(ca_float, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_int, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_str, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_array, ca_dbr_representer, Dumper=yaml.Dumper)
+
     with open(save_path, "w") as file:
         yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)
 

src/ophyd_async/core/flyer.py

@@ -256,10 +256,18 @@ async def _fly(self) -> None:
 
     async def collect_asset_docs(self) -> AsyncIterator[Asset]:
         current_frame = self._current_frame
+        stream_datums: List[Asset] = []
         async for asset in self._detector_group_logic.collect_asset_docs():
             name, doc = asset
             if name == "stream_datum":
                 current_frame = doc["indices"]["stop"] + self._offset
+                # Defer stream_datums until all stream_resources have been produced
+                # In a single collect, all the stream_resources are produced first
+                # followed by their stream_datums
+                stream_datums.append(asset)
+            else:
+                yield asset
+        for asset in stream_datums:
             yield asset
         if current_frame != self._current_frame:
             self._current_frame = current_frame

src/ophyd_async/core/signal.py

@@ -343,7 +343,7 @@ async def set_and_wait_for_value(
 
     Useful for busy record, or other Signals with pattern:
 
-    - Set Signal with wait=with_callback and stash the Status
+    - Set Signal with wait=True and stash the Status
     - Read the same Signal to check the operation has started
     - Return the Status so calling code can wait for operation to complete
 
@@ -357,8 +357,6 @@ async def set_and_wait_for_value(
         The signal to set and monitor
     value:
         The value to set it to
-    with_callback:
-        If we want to wait for a caput/pvput callback
     timeout:
         How long to wait for the signal to have the value
     status_timeout:

src/ophyd_async/epics/areadetector/controllers/ad_sim_controller.py

@@ -1,21 +1,29 @@
 import asyncio
-from typing import Optional
+from typing import Optional, Set
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     AsyncStatus,
     DetectorControl,
     DetectorTrigger,
-    set_and_wait_for_value,
 )
 
-from ..drivers.ad_base import ADBase, ImageMode
+from ..drivers.ad_base import (
+    DEFAULT_GOOD_STATES,
+    ADBase,
+    DetectorState,
+    ImageMode,
+    start_acquiring_driver_and_ensure_status,
+)
 from ..utils import stop_busy_record
 
 
 class ADSimController(DetectorControl):
-    def __init__(self, driver: ADBase) -> None:
+    def __init__(
+        self, driver: ADBase, good_states: Set[DetectorState] = set(DEFAULT_GOOD_STATES)
+    ) -> None:
         self.driver = driver
+        self.good_states = good_states
 
     def get_deadtime(self, exposure: float) -> float:
         return 0.002
@@ -34,8 +42,8 @@ async def arm(
             self.driver.num_images.set(num),
             self.driver.image_mode.set(ImageMode.multiple),
         )
-        return await set_and_wait_for_value(
-            self.driver.acquire, True, timeout=frame_timeout
+        return await start_acquiring_driver_and_ensure_status(
+            self.driver, good_states=self.good_states, timeout=frame_timeout
         )
 
     async def disarm(self):

src/ophyd_async/epics/areadetector/controllers/pilatus_controller.py

@@ -1,11 +1,11 @@
 import asyncio
-from typing import Optional
+from typing import Optional, Set
 
-from ophyd_async.core import (
-    AsyncStatus,
-    DetectorControl,
-    DetectorTrigger,
-    set_and_wait_for_value,
+from ophyd_async.core import AsyncStatus, DetectorControl, DetectorTrigger
+from ophyd_async.epics.areadetector.drivers.ad_base import (
+    DEFAULT_GOOD_STATES,
+    DetectorState,
+    start_acquiring_driver_and_ensure_status,
 )
 
 from ..drivers.pilatus_driver import PilatusDriver, TriggerMode
@@ -19,8 +19,13 @@
 
 
 class PilatusController(DetectorControl):
-    def __init__(self, driver: PilatusDriver) -> None:
+    def __init__(
+        self,
+        driver: PilatusDriver,
+        good_states: Set[DetectorState] = set(DEFAULT_GOOD_STATES),
+    ) -> None:
         self.driver = driver
+        self.good_states = good_states
 
     def get_deadtime(self, exposure: float) -> float:
         return 0.001
@@ -33,10 +38,12 @@ async def arm(
     ) -> AsyncStatus:
         await asyncio.gather(
             self.driver.trigger_mode.set(TRIGGER_MODE[trigger]),
-            self.driver.num_images.set(2**31 - 1 if num == 0 else num),
+            self.driver.num_images.set(999_999 if num == 0 else num),
             self.driver.image_mode.set(ImageMode.multiple),
         )
-        return await set_and_wait_for_value(self.driver.acquire, True)
+        return await start_acquiring_driver_and_ensure_status(
+            self.driver, good_states=self.good_states
+        )
 
     async def disarm(self):
         await stop_busy_record(self.driver.acquire, False, timeout=1)

src/ophyd_async/epics/areadetector/drivers/__init__.py

@@ -1,4 +1,15 @@
-from .ad_base import ADBase, ADBaseShapeProvider
+from .ad_base import (
+    ADBase,
+    ADBaseShapeProvider,
+    DetectorState,
+    start_acquiring_driver_and_ensure_status,
+)
 from .pilatus_driver import PilatusDriver
 
-__all__ = ["ADBase", "ADBaseShapeProvider", "PilatusDriver"]
+__all__ = [
+    "ADBase",
+    "ADBaseShapeProvider",
+    "PilatusDriver",
+    "start_acquiring_driver_and_ensure_status",
+    "DetectorState",
+]

src/ophyd_async/epics/areadetector/drivers/ad_base.py

@@ -1,13 +1,45 @@
 import asyncio
-from typing import Sequence
+from enum import Enum
+from typing import FrozenSet, Sequence, Set
 
-from ophyd_async.core import ShapeProvider
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    AsyncStatus,
+    ShapeProvider,
+    set_and_wait_for_value,
+)
 
 from ...signal.signal import epics_signal_rw
 from ..utils import ImageMode, ad_r, ad_rw
 from ..writers.nd_plugin import NDArrayBase
 
 
+class DetectorState(str, Enum):
+    """
+    Default set of states of an AreaDetector driver.
+    See definition in ADApp/ADSrc/ADDriver.h in https://github.com/areaDetector/ADCore
+    """
+
+    Idle = "Idle"
+    Acquire = "Acquire"
+    Readout = "Readout"
+    Correct = "Correct"
+    Saving = "Saving"
+    Aborting = "Aborting"
+    Error = "Error"
+    Waiting = "Waiting"
+    Initializing = "Initializing"
+    Disconnected = "Disconnected"
+    Aborted = "Aborted"
+
+
+#: Default set of states that we should consider "good" i.e. the acquisition
+#  is complete and went well
+DEFAULT_GOOD_STATES: FrozenSet[DetectorState] = frozenset(
+    [DetectorState.Idle, DetectorState.Aborted]
+)
+
+
 class ADBase(NDArrayBase):
     def __init__(self, prefix: str, name: str = "") -> None:
         # Define some signals
@@ -18,11 +50,55 @@ def __init__(self, prefix: str, name: str = "") -> None:
         self.array_counter = ad_rw(int, prefix + "ArrayCounter")
         self.array_size_x = ad_r(int, prefix + "ArraySizeX")
         self.array_size_y = ad_r(int, prefix + "ArraySizeY")
+        self.detector_state = ad_r(DetectorState, prefix + "DetectorState")
         # There is no _RBV for this one
         self.wait_for_plugins = epics_signal_rw(bool, prefix + "WaitForPlugins")
         super().__init__(prefix, name=name)
 
 
+async def start_acquiring_driver_and_ensure_status(
+    driver: ADBase,
+    good_states: Set[DetectorState] = set(DEFAULT_GOOD_STATES),
+    timeout: float = DEFAULT_TIMEOUT,
+) -> AsyncStatus:
+    """
+    Start acquiring driver, raising ValueError if the detector is in a bad state.
+
+    This sets driver.acquire to True, and waits for it to be True up to a timeout.
+    Then, it checks that the DetectorState PV is in DEFAULT_GOOD_STATES, and otherwise
+    raises a ValueError.
+
+    Parameters
+    ----------
+    driver:
+        The driver to start acquiring. Must subclass ADBase.
+    good_states:
+        set of states defined in DetectorState enum which are considered good states.
+    timeout:
+        How long to wait for driver.acquire to readback True (i.e. acquiring).
+
+    Returns
+    -------
+    AsyncStatus:
+        An AsyncStatus that can be awaited to set driver.acquire to True and perform
+        subsequent raising (if applicable) due to detector state.
+    """
+
+    status = await set_and_wait_for_value(driver.acquire, True, timeout=timeout)
+
+    async def complete_acquisition() -> None:
+        """NOTE: possible race condition here between the callback from
+        set_and_wait_for_value and the detector state updating."""
+        await status
+        state = await driver.detector_state.get_value()
+        if state not in good_states:
+            raise ValueError(
+                f"Final detector state {state} not in valid end states: {good_states}"
+            )
+
+    return AsyncStatus(complete_acquisition())
+
+
 class ADBaseShapeProvider(ShapeProvider):
     def __init__(self, driver: ADBase) -> None:
         self._driver = driver