Plan wrapper functions to allow for easier cleanup of DAE setup changes

doc/plan_stubs/plan_wrappers.md

@@ -0,0 +1,79 @@
+# Plan Wrappers
+
+Plan wrappers that temporarily modify [DAE (Data Acquisition Electronics)](https://isiscomputinggroup.github.io/ibex_bluesky_core/devices/dae.html#dae-data-acquisition-electronics) settings during a plan, automatically restoring the original values afterwards. This ensures that your experiments don't permanently change instrument configuration.
+
+## Available Wrappers
+
+### DAE Table
+
+:py:obj:`dae_table_wrapper <ibex_bluesky_core.plan_stubs.dae_table_wrapper>`
+
+```python
+RE(
+    _with_dae_tables(
+        bps.null(),
+        dae=dae,
+        new_settings=modified_settings
+        )
+    )
+```
+
+Where `modified_settings` is a dataset in the form :py:obj:`DaeSettingsData < ibex_bluesky_core.devices.dae.DaeSettingsData>`
+
+A function that wraps a plan to temporarily modify the DAE table.
+
+### Num Periods
+:py:obj:`num_periods_wrapper <ibex_bluesky_core.plan_stubs.num_periods_wrapper>`
+
+```python
+RE(
+    _with_num_periods(
+        bps.null(),
+        dae=dae,
+        number_of_periods=1000 
+        )
+    )
+```
+A function that wraps a plan to temporarily modify the number of periods.
+
+### Time Channels
+:py:obj:`time_channels_wrapper <ibex_bluesky_core.plan_stubs.time_channels_wrapper>`:
+
+```python
+RE(
+    _with_time_channels(
+        bps.null(),
+        dae=dae,
+        new_settings=modified_settings
+        )
+    )
+```
+Where `modified_settings` is a dataset in the form :py:obj:`DaeTCBSettingsData < ibex_bluesky_core.devices.dae.DaeTCBSettingsData>`
+
+A function that wraps a plan to temporarily modify the time channels boundaries.
+
+## Usage
+
+To use these wrappers, the plan written by the user must be wrapped by the function within the RunEngine:
+
+``` python
+
+from bluesky import RunEngine
+from ibex_bluesky_core.plan_stubs import _with_num_periods
+from ibex_bluesky_core.devices.simpledae import SimpleDae
+
+dae = SimpleDae() # Give your DAE options here
+RE = RunEngine()
+
+RE(
+    _with_num_periods(
+        bps.null(), # Default plan to run
+        dae=dae,
+        number_of_periods=1000 # Temporary number of periods to run
+        )
+    )
+
+```
+
+the plan with the modified DAE settings, restoring the original settings afterwards.
+

src/ibex_bluesky_core/devices/dae/_period_settings.py

@@ -7,7 +7,7 @@
 from xml.etree.ElementTree import tostring
 
 from bluesky.protocols import Locatable, Location, Movable
-from ophyd_async.core import AsyncStatus, Device, SignalRW
+from ophyd_async.core import AsyncStatus, SignalRW, StandardReadable
 
 from ibex_bluesky_core.devices import (
     isis_epics_signal_rw,
@@ -108,7 +108,9 @@ def _convert_period_settings_to_xml(current_xml: str, value: DaePeriodSettingsDa
     return tostring(root, encoding="unicode")
 
 
-class DaePeriodSettings(Device, Locatable[DaePeriodSettingsData], Movable[DaePeriodSettingsData]):
+class DaePeriodSettings(
+    StandardReadable, Locatable[DaePeriodSettingsData], Movable[DaePeriodSettingsData]
+):
     """Subdevice for the DAE hardware period settings."""
 
     def __init__(self, dae_prefix: str, name: str = "") -> None:

src/ibex_bluesky_core/devices/dae/_tcb_settings.py

@@ -7,7 +7,7 @@
 from xml.etree.ElementTree import tostring
 
 from bluesky.protocols import Locatable, Location, Movable
-from ophyd_async.core import AsyncStatus, Device, SignalRW
+from ophyd_async.core import AsyncStatus, SignalRW, StandardReadable
 
 from ibex_bluesky_core.devices import (
     compress_and_hex,
@@ -120,7 +120,7 @@ def _convert_tcb_settings_to_xml(current_xml: str, settings: DaeTCBSettingsData)
     return tostring(root, encoding="unicode")
 
 
-class DaeTCBSettings(Device, Locatable[DaeTCBSettingsData], Movable[DaeTCBSettingsData]):
+class DaeTCBSettings(StandardReadable, Locatable[DaeTCBSettingsData], Movable[DaeTCBSettingsData]):
     """Subdevice for the DAE time channel settings."""
 
     def __init__(self, dae_prefix: str, name: str = "") -> None:

src/ibex_bluesky_core/plan_stubs.py → src/ibex_bluesky_core/plan_stubs/__init__.py

@@ -12,6 +12,9 @@
 from ophyd_async.epics.motor import Motor, UseSetMode
 
 from ibex_bluesky_core.devices.reflectometry import ReflParameter
+from ibex_bluesky_core.plan_stubs.dae_table_wrapper import _with_dae_tables
+from ibex_bluesky_core.plan_stubs.num_periods_wrapper import _with_num_periods
+from ibex_bluesky_core.plan_stubs.time_channels_wrapper import _with_time_channels
 from ibex_bluesky_core.utils import NamedReadableAndMovable
 
 logger = logging.getLogger(__name__)
@@ -25,6 +28,9 @@
 
 
 __all__ = [
+    "_with_dae_tables",
+    "_with_num_periods",
+    "_with_time_channels",
     "call_qt_aware",
     "call_sync",
     "polling_plan",

src/ibex_bluesky_core/plan_stubs/dae_table_wrapper.py

@@ -0,0 +1,43 @@
+"""Wrap a plan with temporary modification to DAE Settings."""
+
+from collections.abc import Generator
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import Msg
+from ophyd_async.plan_stubs import ensure_connected
+
+from ibex_bluesky_core.devices.dae import Dae, DaeSettingsData
+
+
+def _with_dae_tables(
+    plan: Generator[Msg, None, None], dae: Dae, new_settings: DaeSettingsData
+) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to DAE Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+        new_settings: The new DAE Settings to apply temporarily.
+
+    Returns:
+        A generator which runs the plan with the modified DAE settings, restoring the original
+        settings afterwards.
+
+    """
+    yield from ensure_connected(dae)
+
+    original_dae_setting = None
+
+    def _inner() -> Generator[Msg, None, None]:
+        nonlocal original_dae_setting
+        original_dae_setting = yield from bps.rd(dae.dae_settings)
+
+        yield from bps.mv(dae.dae_settings, new_settings)
+
+        yield from plan
+
+    def _cleanup() -> Generator[Msg, None, None]:
+        yield from bps.mv(dae.dae_settings, original_dae_setting)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _cleanup()))

src/ibex_bluesky_core/plan_stubs/num_periods_wrapper.py

@@ -0,0 +1,42 @@
+"""Wrap a plan with temporary modification to Periods Settings."""
+
+from collections.abc import Generator
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import Msg
+from ophyd_async.plan_stubs import ensure_connected
+
+from ibex_bluesky_core.devices.dae import Dae
+
+
+def _with_num_periods(
+    plan: Generator[Msg, None, None], dae: Dae, number_of_periods: int
+) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to Periods Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+        number_of_periods: The number of periods to set to temporarily.
+
+    Returns:
+        A generator which runs the plan with the modified number of periods, restoring the original
+        number of periods afterwards.
+
+    """
+    original_num_periods = None
+
+    def _inner() -> Generator[Msg, None, None]:
+        yield from ensure_connected(dae)
+        nonlocal original_num_periods
+        original_num_periods = yield from bps.rd(dae.number_of_periods)
+
+        yield from bps.mv(dae.number_of_periods, number_of_periods)
+
+        yield from plan
+
+    def _cleanup() -> Generator[Msg, None, None]:
+        yield from bps.mv(dae.number_of_periods, original_num_periods)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _cleanup()))

src/ibex_bluesky_core/plan_stubs/time_channels_wrapper.py

@@ -0,0 +1,43 @@
+"""Wrap a plan with temporary modification to Time Channel Settings."""
+
+from collections.abc import Generator
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import Msg
+from ophyd_async.plan_stubs import ensure_connected
+
+from ibex_bluesky_core.devices.dae import Dae, DaeTCBSettingsData
+
+
+def _with_time_channels(
+    plan: Generator[Msg, None, None], dae: Dae, new_tcb_settings: DaeTCBSettingsData
+) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to Time Channel Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+        new_tcb_settings: The time channel settings to apply temporarily.
+
+    Returns:
+        A generator which runs the plan with the modified TCB settings, restoring the original
+        settings afterwards.
+
+    """
+    yield from ensure_connected(dae)
+
+    original_time_channels = None
+
+    def _inner() -> Generator[Msg, None, None]:
+        nonlocal original_time_channels
+        original_time_channels = yield from bps.rd(dae.tcb_settings)
+
+        yield from bps.mv(dae.tcb_settings, new_tcb_settings)
+
+        yield from plan
+
+    def _cleanup() -> Generator[Msg, None, None]:
+        yield from bps.mv(dae.tcb_settings, original_time_channels)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _cleanup()))

tests/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 from bluesky.run_engine import RunEngine
 
+from ibex_bluesky_core.devices.dae import Dae
 from ibex_bluesky_core.devices.simpledae import Controller, Reducer, SimpleDae, Waiter
 from ibex_bluesky_core.run_engine import get_run_engine
 
@@ -36,3 +37,10 @@ async def simpledae() -> SimpleDae:
     )
     await dae.connect(mock=True)
     return dae
+
+
+@pytest.fixture
+async def dae() -> Dae:
+    dae = Dae("UNITTEST:MOCK:")
+    await dae.connect(mock=True)
+    return dae

tests/devices/test_dae.py

@@ -50,13 +50,6 @@
 )
 
 
-@pytest.fixture
-async def dae() -> Dae:
-    dae = Dae("UNITTEST:MOCK:")
-    await dae.connect(mock=True)
-    return dae
-
-
 @pytest.fixture
 async def spectrum() -> DaeSpectra:
     spectrum = DaeSpectra(dae_prefix="UNITTEST:MOCK:", spectra=1, period=1)