Source code for ophyd_async.epics.adcore._acquire_logic

import asyncio

from ophyd_async.core import (
    DEFAULT_TIMEOUT,
    AsyncStatus,
    DetectorAcquireLogic,
    SignalR,
    set_and_wait_for_other_value,
    set_and_wait_for_value,
    wait_for_value,
)
from ophyd_async.epics.core import stop_busy_record, wait_for_good_state

from ._io import ADBaseIO, ADState, NDCircularBuffIO




class ADAcquireLogic(DetectorAcquireLogic):
    def __init__(
        self,
        driver: ADBaseIO,
        driver_armed_signal: SignalR[bool] | None = None,
        timeout: float = DEFAULT_TIMEOUT,
    ):
        self.driver = driver
        if driver_armed_signal is not None:
            self.driver_armed_signal = driver_armed_signal
        else:
            self.driver_armed_signal = driver.acquire
        self.timeout = timeout
        self.acquire_status: AsyncStatus | None = None



    async def start_acquiring(self):
        self.acquire_status = await set_and_wait_for_other_value(
            set_signal=self.driver.acquire,
            set_value=True,
            match_signal=self.driver_armed_signal,
            match_value=True,
            wait_for_set_completion=False,
            timeout=self.timeout,
        )




    async def wait_for_idle(self):
        if self.acquire_status:
            await self.acquire_status
        await wait_for_good_state(
            self.driver.detector_state,
            {ADState.IDLE, ADState.ABORTED},
            timeout=self.timeout,
        )




    async def ensure_stopped(self):
        await stop_busy_record(self.driver.acquire)






class ADContAcqAcquireLogic(DetectorAcquireLogic):
    def __init__(self, driver: ADBaseIO, cb_plugin: NDCircularBuffIO):
        self.driver = driver
        self.cb_plugin = cb_plugin
        self.acquire_status: AsyncStatus | None = None



    async def ensure_ready(self):
        await set_and_wait_for_value(
            self.cb_plugin.capture,
            True,
            wait_for_set_completion=False,
            timeout=DEFAULT_TIMEOUT,
        )




    async def start_acquiring(self):
        await self.cb_plugin.trigger_.set(True)




    async def wait_for_idle(self):
        num_images, acquire_time = await asyncio.gather(
            self.cb_plugin.post_count.get_value(),
            self.driver.acquire_time.get_value(),
        )

        await wait_for_value(
            self.cb_plugin.trigger_,
            False,
            timeout=num_images * acquire_time + DEFAULT_TIMEOUT,
        )




    async def ensure_stopped(self):
        await stop_busy_record(self.cb_plugin.capture)