ophyd-async 0.8.0a5__py3-none-any.whl → 0.9.0__py3-none-any.whl

ophyd_async/epics/adcore/_core_io.py

@@ -1,16 +1,19 @@
+import asyncio
+
 from ophyd_async.core import Device, StrictEnum
+from ophyd_async.core._providers import DatasetDescriber
 from ophyd_async.epics.core import (
     epics_signal_r,
     epics_signal_rw,
     epics_signal_rw_rbv,
 )
 
-from ._utils import ADBaseDataType, FileWriteMode, ImageMode
+from ._utils import ADBaseDataType, FileWriteMode, ImageMode, convert_ad_dtype_to_np
 
 
 class Callback(StrictEnum):
-    Enable = "Enable"
-    Disable = "Disable"
+    ENABLE = "Enable"
+    DISABLE = "Disable"
 
 
 class NDArrayBaseIO(Device):
@@ -27,6 +30,21 @@ class NDArrayBaseIO(Device):
         super().__init__(name=name)
 
 
+class ADBaseDatasetDescriber(DatasetDescriber):
+    def __init__(self, driver: NDArrayBaseIO) -> None:
+        self._driver = driver
+
+    async def np_datatype(self) -> str:
+        return convert_ad_dtype_to_np(await self._driver.data_type.get_value())
+
+    async def shape(self) -> tuple[int, int]:
+        shape = await asyncio.gather(
+            self._driver.array_size_y.get_value(),
+            self._driver.array_size_x.get_value(),
+        )
+        return shape
+
+
 class NDPluginBaseIO(NDArrayBaseIO):
     def __init__(self, prefix: str, name: str = "") -> None:
         self.nd_array_port = epics_signal_rw_rbv(str, prefix + "NDArrayPort")
@@ -36,6 +54,7 @@ class NDPluginBaseIO(NDArrayBaseIO):
         self.nd_array_address = epics_signal_rw_rbv(int, prefix + "NDArrayAddress")
         self.array_size0 = epics_signal_r(int, prefix + "ArraySize0_RBV")
         self.array_size1 = epics_signal_r(int, prefix + "ArraySize1_RBV")
+        self.queue_size = epics_signal_rw(int, prefix + "QueueSize")
         super().__init__(prefix, name)
 
 
@@ -72,17 +91,17 @@ class DetectorState(StrictEnum):
     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"
+    IDLE = "Idle"
+    ACQUIRE = "Acquire"
+    READOUT = "Readout"
+    CORRECT = "Correct"
+    SAVING = "Saving"
+    ABORTING = "Aborting"
+    ERROR = "Error"
+    WAITING = "Waiting"
+    INITIALIZING = "Initializing"
+    DISCONNECTED = "Disconnected"
+    ABORTED = "Aborted"
 
 
 class ADBaseIO(NDArrayBaseIO):
@@ -99,40 +118,46 @@ class ADBaseIO(NDArrayBaseIO):
 
 
 class Compression(StrictEnum):
-    none = "None"
-    nbit = "N-bit"
-    szip = "szip"
-    zlib = "zlib"
-    blosc = "Blosc"
-    bslz4 = "BSLZ4"
-    lz4 = "LZ4"
-    jpeg = "JPEG"
+    NONE = "None"
+    NBIT = "N-bit"
+    SZIP = "szip"
+    ZLIB = "zlib"
+    BLOSC = "Blosc"
+    BSLZ4 = "BSLZ4"
+    LZ4 = "LZ4"
+    JPEG = "JPEG"
 
 
-class NDFileHDFIO(NDPluginBaseIO):
+class NDFileIO(NDPluginBaseIO):
     def __init__(self, prefix: str, name="") -> None:
-        # Define some signals
-        self.position_mode = epics_signal_rw_rbv(bool, prefix + "PositionMode")
-        self.compression = epics_signal_rw_rbv(Compression, prefix + "Compression")
-        self.num_extra_dims = epics_signal_rw_rbv(int, prefix + "NumExtraDims")
         self.file_path = epics_signal_rw_rbv(str, prefix + "FilePath")
         self.file_name = epics_signal_rw_rbv(str, prefix + "FileName")
         self.file_path_exists = epics_signal_r(bool, prefix + "FilePathExists_RBV")
         self.file_template = epics_signal_rw_rbv(str, prefix + "FileTemplate")
         self.full_file_name = epics_signal_r(str, prefix + "FullFileName_RBV")
+        self.file_number = epics_signal_rw(int, prefix + "FileNumber")
+        self.auto_increment = epics_signal_rw(bool, prefix + "AutoIncrement")
         self.file_write_mode = epics_signal_rw_rbv(
             FileWriteMode, prefix + "FileWriteMode"
         )
         self.num_capture = epics_signal_rw_rbv(int, prefix + "NumCapture")
         self.num_captured = epics_signal_r(int, prefix + "NumCaptured_RBV")
-        self.swmr_mode = epics_signal_rw_rbv(bool, prefix + "SWMRMode")
-        self.lazy_open = epics_signal_rw_rbv(bool, prefix + "LazyOpen")
         self.capture = epics_signal_rw_rbv(bool, prefix + "Capture")
-        self.flush_now = epics_signal_rw(bool, prefix + "FlushNow")
-        self.xml_file_name = epics_signal_rw_rbv(str, prefix + "XMLFileName")
         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")
+        super().__init__(prefix, name)
+
+
+class NDFileHDFIO(NDFileIO):
+    def __init__(self, prefix: str, name="") -> None:
+        self.position_mode = epics_signal_rw_rbv(bool, prefix + "PositionMode")
+        self.compression = epics_signal_rw_rbv(Compression, prefix + "Compression")
+        self.num_extra_dims = epics_signal_rw_rbv(int, prefix + "NumExtraDims")
+        self.swmr_mode = epics_signal_rw_rbv(bool, prefix + "SWMRMode")
+        self.flush_now = epics_signal_rw(bool, prefix + "FlushNow")
+        self.xml_file_name = epics_signal_rw_rbv(str, prefix + "XMLFileName")
         self.num_frames_chunks = epics_signal_r(int, prefix + "NumFramesChunks_RBV")
         self.chunk_size_auto = epics_signal_rw_rbv(bool, prefix + "ChunkSizeAuto")
+        self.lazy_open = epics_signal_rw_rbv(bool, prefix + "LazyOpen")
         super().__init__(prefix, name)

ophyd_async/epics/adcore/_core_logic.py

@@ -1,106 +1,121 @@
 import asyncio
+from typing import Generic, TypeVar
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     AsyncStatus,
-    DatasetDescriber,
     DetectorController,
+    DetectorTrigger,
+    TriggerInfo,
     set_and_wait_for_value,
 )
-from ophyd_async.epics.adcore._utils import convert_ad_dtype_to_np
 
 from ._core_io import ADBaseIO, DetectorState
+from ._utils import ImageMode, stop_busy_record
 
 # 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]
+    [DetectorState.IDLE, DetectorState.ABORTED]
 )
 
-
-class ADBaseDatasetDescriber(DatasetDescriber):
-    def __init__(self, driver: ADBaseIO) -> None:
-        self._driver = driver
-
-    async def np_datatype(self) -> str:
-        return convert_ad_dtype_to_np(await self._driver.data_type.get_value())
-
-    async def shape(self) -> tuple[int, int]:
-        shape = await asyncio.gather(
-            self._driver.array_size_y.get_value(),
-            self._driver.array_size_x.get_value(),
+ADBaseIOT = TypeVar("ADBaseIOT", bound=ADBaseIO)
+ADBaseControllerT = TypeVar("ADBaseControllerT", bound="ADBaseController")
+
+
+class ADBaseController(DetectorController, Generic[ADBaseIOT]):
+    def __init__(
+        self,
+        driver: ADBaseIOT,
+        good_states: frozenset[DetectorState] = DEFAULT_GOOD_STATES,
+    ) -> None:
+        self.driver = driver
+        self.good_states = good_states
+        self.frame_timeout = DEFAULT_TIMEOUT
+        self._arm_status: AsyncStatus | None = None
+
+    async def prepare(self, trigger_info: TriggerInfo) -> None:
+        if trigger_info.trigger != DetectorTrigger.INTERNAL:
+            msg = (
+                "fly scanning (i.e. external triggering) is not supported for this "
+                "device"
+            )
+            raise TypeError(msg)
+        self.frame_timeout = (
+            DEFAULT_TIMEOUT + await self.driver.acquire_time.get_value()
         )
-        return shape
-
-
-async def set_exposure_time_and_acquire_period_if_supplied(
-    controller: DetectorController,
-    driver: ADBaseIO,
-    exposure: float | None = None,
-    timeout: float = DEFAULT_TIMEOUT,
-) -> None:
-    """
-    Sets the exposure time if it is not None and the acquire period to the
-    exposure time plus the deadtime. This is expected behavior for most
-    AreaDetectors, but some may require more specialized handling.
-
-    Parameters
-    ----------
-    controller:
-        Controller that can supply a deadtime.
-    driver:
-        The driver to start acquiring. Must subclass ADBaseIO.
-    exposure:
-        Desired exposure time, this is a noop if it is None.
-    timeout:
-        How long to wait for the exposure time and acquire period to be set.
-    """
-    if exposure is not None:
-        full_frame_time = exposure + controller.get_deadtime(exposure)
         await asyncio.gather(
-            driver.acquire_time.set(exposure, timeout=timeout),
-            driver.acquire_period.set(full_frame_time, timeout=timeout),
+            self.driver.num_images.set(trigger_info.total_number_of_triggers),
+            self.driver.image_mode.set(ImageMode.MULTIPLE),
         )
 
-
-async def start_acquiring_driver_and_ensure_status(
-    driver: ADBaseIO,
-    good_states: frozenset[DetectorState] = frozenset(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 ADBaseIO.
-    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}"
+    async def arm(self):
+        self._arm_status = await self.start_acquiring_driver_and_ensure_status()
+
+    async def wait_for_idle(self):
+        if self._arm_status:
+            await self._arm_status
+
+    async def disarm(self):
+        # We can't use caput callback as we already used it in arm() and we can't have
+        # 2 or they will deadlock
+        await stop_busy_record(self.driver.acquire, False, timeout=1)
+
+    async def set_exposure_time_and_acquire_period_if_supplied(
+        self,
+        exposure: float | None = None,
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> None:
+        """
+        Sets the exposure time if it is not None and the acquire period to the
+        exposure time plus the deadtime. This is expected behavior for most
+        AreaDetectors, but some may require more specialized handling.
+
+        Parameters
+        ----------
+        exposure:
+            Desired exposure time, this is a noop if it is None.
+        timeout:
+            How long to wait for the exposure time and acquire period to be set.
+        """
+        if exposure is not None:
+            full_frame_time = exposure + self.get_deadtime(exposure)
+            await asyncio.gather(
+                self.driver.acquire_time.set(exposure, timeout=timeout),
+                self.driver.acquire_period.set(full_frame_time, timeout=timeout),
             )
 
-    return AsyncStatus(complete_acquisition())
+    async def start_acquiring_driver_and_ensure_status(self) -> 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.
+
+        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(
+            self.driver.acquire,
+            True,
+            timeout=DEFAULT_TIMEOUT,
+            wait_for_set_completion=False,
+        )
+
+        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 self.driver.detector_state.get_value()
+            if state not in self.good_states:
+                raise ValueError(
+                    f"Final detector state {state.value} not "
+                    "in valid end states: {self.good_states}"
+                )
+
+        return AsyncStatus(complete_acquisition())

ophyd_async/epics/adcore/_core_writer.py

@@ -0,0 +1,219 @@
+import asyncio
+from collections.abc import AsyncGenerator, AsyncIterator
+from pathlib import Path
+from typing import Generic, TypeVar, get_args
+from urllib.parse import urlunparse
+
+from bluesky.protocols import Hints, StreamAsset
+from event_model import (
+    ComposeStreamResource,
+    DataKey,
+    StreamRange,
+)
+
+from ophyd_async.core._detector import DetectorWriter
+from ophyd_async.core._providers import DatasetDescriber, NameProvider, PathProvider
+from ophyd_async.core._signal import (
+    observe_value,
+    set_and_wait_for_value,
+    wait_for_value,
+)
+from ophyd_async.core._status import AsyncStatus
+from ophyd_async.core._utils import DEFAULT_TIMEOUT
+
+# from ophyd_async.epics.adcore._core_logic import ADBaseDatasetDescriber
+from ._core_io import (
+    ADBaseDatasetDescriber,
+    Callback,
+    NDArrayBaseIO,
+    NDFileIO,
+    NDPluginBaseIO,
+)
+from ._utils import FileWriteMode
+
+NDFileIOT = TypeVar("NDFileIOT", bound=NDFileIO)
+ADWriterT = TypeVar("ADWriterT", bound="ADWriter")
+
+
+class ADWriter(DetectorWriter, Generic[NDFileIOT]):
+    default_suffix: str = "FILE1:"
+
+    def __init__(
+        self,
+        fileio: NDFileIOT,
+        path_provider: PathProvider,
+        name_provider: NameProvider,
+        dataset_describer: DatasetDescriber,
+        file_extension: str = "",
+        mimetype: str = "",
+        plugins: dict[str, NDPluginBaseIO] | None = None,
+    ) -> None:
+        self._plugins = plugins or {}
+        self.fileio = fileio
+        self._path_provider = path_provider
+        self._name_provider = name_provider
+        self._dataset_describer = dataset_describer
+        self._file_extension = file_extension
+        self._mimetype = mimetype
+        self._last_emitted = 0
+        self._emitted_resource = None
+
+        self._capture_status: AsyncStatus | None = None
+        self._multiplier = 1
+        self._filename_template = "%s%s_%6.6d"
+
+    @classmethod
+    def with_io(
+        cls: type[ADWriterT],
+        prefix: str,
+        path_provider: PathProvider,
+        dataset_source: NDArrayBaseIO | None = None,
+        fileio_suffix: str | None = None,
+        plugins: dict[str, NDPluginBaseIO] | None = None,
+    ) -> ADWriterT:
+        try:
+            fileio_cls = get_args(cls.__orig_bases__[0])[0]  # type: ignore
+        except IndexError as err:
+            raise RuntimeError("File IO class for writer not specified!") from err
+
+        fileio = fileio_cls(prefix + (fileio_suffix or cls.default_suffix))
+        dataset_describer = ADBaseDatasetDescriber(dataset_source or fileio)
+
+        def name_provider() -> str:
+            if fileio.parent == "Not attached to a detector":
+                raise RuntimeError("Initializing writer without parent detector!")
+            return fileio.parent.name
+
+        writer = cls(
+            fileio, path_provider, name_provider, dataset_describer, plugins=plugins
+        )
+        return writer
+
+    async def begin_capture(self) -> None:
+        info = self._path_provider(device_name=self._name_provider())
+
+        await self.fileio.enable_callbacks.set(Callback.ENABLE)
+
+        # Set the directory creation depth first, since dir creation callback happens
+        # when directory path PV is processed.
+        await self.fileio.create_directory.set(info.create_dir_depth)
+
+        await asyncio.gather(
+            # See https://github.com/bluesky/ophyd-async/issues/122
+            self.fileio.file_path.set(str(info.directory_path)),
+            self.fileio.file_name.set(info.filename),
+            self.fileio.file_write_mode.set(FileWriteMode.STREAM),
+            # For non-HDF file writers, use AD file templating mechanism
+            # for generating multi-image datasets
+            self.fileio.file_template.set(
+                self._filename_template + self._file_extension
+            ),
+            self.fileio.auto_increment.set(True),
+            self.fileio.file_number.set(0),
+        )
+
+        if not await self.fileio.file_path_exists.get_value():
+            msg = f"File path {info.directory_path} for file plugin does not exist"
+            raise FileNotFoundError(msg)
+
+        # Overwrite num_capture to go forever
+        await self.fileio.num_capture.set(0)
+        # Wait for it to start, stashing the status that tells us when it finishes
+        self._capture_status = await set_and_wait_for_value(
+            self.fileio.capture, True, wait_for_set_completion=False
+        )
+
+    async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
+        self._emitted_resource = None
+        self._last_emitted = 0
+        self._multiplier = multiplier
+        frame_shape = await self._dataset_describer.shape()
+        dtype_numpy = await self._dataset_describer.np_datatype()
+
+        await self.begin_capture()
+
+        describe = {
+            self._name_provider(): DataKey(
+                source=self._name_provider(),
+                shape=list(frame_shape),
+                dtype="array",
+                dtype_numpy=dtype_numpy,
+                external="STREAM:",
+            )  # type: ignore
+        }
+        return describe
+
+    async def observe_indices_written(
+        self, timeout=DEFAULT_TIMEOUT
+    ) -> AsyncGenerator[int, None]:
+        """Wait until a specific index is ready to be collected"""
+        async for num_captured in observe_value(self.fileio.num_captured, timeout):
+            yield num_captured // self._multiplier
+
+    async def get_indices_written(self) -> int:
+        num_captured = await self.fileio.num_captured.get_value()
+        return num_captured // self._multiplier
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        if indices_written:
+            if not self._emitted_resource:
+                file_path = Path(await self.fileio.file_path.get_value())
+                file_name = await self.fileio.file_name.get_value()
+                file_template = file_name + "_{:06d}" + self._file_extension
+
+                frame_shape = await self._dataset_describer.shape()
+
+                uri = urlunparse(
+                    (
+                        "file",
+                        "localhost",
+                        str(file_path.absolute()) + "/",
+                        "",
+                        "",
+                        None,
+                    )
+                )
+
+                bundler_composer = ComposeStreamResource()
+
+                self._emitted_resource = bundler_composer(
+                    mimetype=self._mimetype,
+                    uri=uri,
+                    data_key=self._name_provider(),
+                    parameters={
+                        # Assume that we always write 1 frame per file/chunk
+                        "chunk_shape": (1, *frame_shape),
+                        # Include file template for reconstruction in consolidator
+                        "template": file_template,
+                    },
+                    uid=None,
+                    validate=True,
+                )
+
+                yield "stream_resource", self._emitted_resource.stream_resource_doc
+
+            # Indices are relative to resource
+            if indices_written > self._last_emitted:
+                indices: StreamRange = {
+                    "start": self._last_emitted,
+                    "stop": indices_written,
+                }
+                self._last_emitted = indices_written
+                yield (
+                    "stream_datum",
+                    self._emitted_resource.compose_stream_datum(indices),
+                )
+
+    async def close(self):
+        # Already done a caput callback in _capture_status, so can't do one here
+        await self.fileio.capture.set(False, wait=False)
+        await wait_for_value(self.fileio.capture, False, DEFAULT_TIMEOUT)
+        if self._capture_status:
+            # We kicked off an open, so wait for it to return
+            await self._capture_status
+
+    @property
+    def hints(self) -> Hints:
+        return {"fields": [self._name_provider()]}