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

doc/plan_stubs/plan_wrappers.md

@@ -0,0 +1,39 @@
+# Plan Wrappers
+
+Plan wrappers that temporarily modify 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_wrapper](ibex_bluesky_core.plan_stubs.dae_table_wrapper)
+['num_periods_wrapper](ibex_bluesky_core.plan_stubs.num_periods_wrapper)
+['time_channels_wrapper](ibex_bluesky_core.plan_stubs.time_channels_wrapper)
+
+## 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()
+
+modified_settings = 1
+
+def example_plan():
+    yield from bps.mv(dae.number_of_periods, modified_settings)
+
+RE(
+    with_num_periods(
+        example_plan(), 
+        dae=dae,
+        )
+    )
+
+```
+
+the plan with the modified DAE settings, restoring the original settings afterwards.
+

src/ibex_bluesky_core/devices/dae/_period_settings.py

@@ -53,6 +53,14 @@ class SinglePeriodSettings:
     label: str | None = None
 
 
+class PeriodSettingType(Enum):
+    """Periods type option for a single row."""
+
+    UNUSED = 0
+    DAQ = 1
+    DWEll = 2
+
+
 @dataclass(kw_only=True)
 class DaePeriodSettingsData:
     """Dataclass for the hardware period settings."""

src/ibex_bluesky_core/plan_stubs/dae_table_wrapper.py

@@ -0,0 +1,39 @@
+"""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
+
+
+def with_dae_tables(plan: Generator[Msg, None, None], dae: Dae) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to DAE Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+
+    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 plan
+
+    def _onexit() -> Generator[Msg, None, None]:
+        nonlocal original_dae_setting
+        yield from bps.mv(dae.dae_settings, original_dae_setting)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _onexit()))

src/ibex_bluesky_core/plan_stubs/num_periods_wrapper.py

@@ -0,0 +1,39 @@
+"""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) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to Periods Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+
+    Returns:
+        A generator which runs the plan with the modified DAE settings, restoring the original
+        settings afterwards.
+
+    """
+    yield from ensure_connected(dae)
+
+    original_num_periods = None
+
+    def _inner() -> Generator[Msg, None, None]:
+        nonlocal original_num_periods
+        original_num_periods = yield from bps.rd(dae.number_of_periods)
+
+        yield from plan
+
+    def _onexit() -> Generator[Msg, None, None]:
+        nonlocal original_num_periods
+        yield from bps.mv(dae.number_of_periods, original_num_periods)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _onexit()))

src/ibex_bluesky_core/plan_stubs/time_channels_wrapper.py

@@ -0,0 +1,39 @@
+"""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
+
+
+def with_time_channels(plan: Generator[Msg, None, None], dae: Dae) -> Generator[Msg, None, None]:
+    """Wrap a plan with temporary modification to Time Channel Settings.
+
+    Args:
+        plan: The plan to wrap.
+        dae: The Dae instance.
+
+    Returns:
+        A generator which runs the plan with the modified DAE 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)  # type: ignore
+
+        yield from plan
+
+    def _onexit() -> Generator[Msg, None, None]:
+        nonlocal original_time_channels
+        yield from bps.mv(dae.tcb_settings, original_time_channels)
+
+    return (yield from bpp.finalize_wrapper(_inner(), _onexit()))

src/ibex_bluesky_core/plans/__init__.py

@@ -16,6 +16,7 @@
 from ibex_bluesky_core.devices.simpledae import monitor_normalising_dae
 from ibex_bluesky_core.fitting import FitMethod
 from ibex_bluesky_core.plan_stubs import call_qt_aware, polling_plan
+from ibex_bluesky_core.plan_stubs.num_periods_wrapper import with_num_periods
 from ibex_bluesky_core.utils import NamedReadableAndMovable, centred_pixel
 
 if TYPE_CHECKING:
@@ -28,6 +29,7 @@
     "motor_scan",
     "polling_plan",
     "scan",
+    "with_num_periods",
 ]
 
 

tests/devices/test_dae.py

@@ -1,6 +1,6 @@
 # pyright: reportMissingParameterType=false
 from enum import Enum
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 from xml.etree import ElementTree as ET
 
 import bluesky.plan_stubs as bps
@@ -39,6 +39,9 @@
 )
 from ibex_bluesky_core.devices.dae._period_settings import _convert_period_settings_to_xml
 from ibex_bluesky_core.devices.dae._tcb_settings import _convert_tcb_settings_to_xml
+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 tests.conftest import MOCK_PREFIX
 from tests.devices.dae_testing_data import (
     dae_settings_template,
@@ -1023,3 +1026,96 @@ async def test_if_tof_edges_has_no_units_then_read_spec_dataarray_gives_error(
 
 def test_dae_repr():
     assert repr(Dae(prefix="foo", name="bar")) == "Dae(name=bar, prefix=foo)"
+
+
+async def test_num_periods_wrapper(dae: Dae, RE: RunEngine):
+    original_settings = 4
+    modified_settings = 7
+
+    await dae.number_of_periods.set(original_settings)
+
+    def _dummy_plan_which_sets_periods(dae):
+        yield from bps.mv(dae.number_of_periods, modified_settings)
+
+        current = yield from bps.rd(dae.number_of_periods)
+        assert current == 7
+
+    with patch("ibex_bluesky_core.plan_stubs.num_periods_wrapper.ensure_connected"):
+        RE(
+            with_num_periods(
+                _dummy_plan_which_sets_periods(dae),
+                dae=dae,  # type: ignore
+            )
+        )
+
+    result = await dae.number_of_periods.read()
+    assert result["DAE-number_of_periods-signal"]["value"] == original_settings
+
+
+def test_time_channels_wrapper(dae: Dae, RE: RunEngine):
+    original_settings = compress_and_hex(initial_tcb_settings).decode()
+    modified_settings = DaeTCBSettingsData(time_unit=TCBTimeUnit.NANOSECONDS)
+
+    expected_time_units = TCBTimeUnit.MICROSECONDS
+
+    set_mock_value(dae.tcb_settings._raw_tcb_settings, original_settings)
+
+    before: DaeTCBSettingsData = RE(bps.rd(dae.tcb_settings)).plan_result  # type: ignore
+
+    def _dummy_plan_which_sets_time_units(dae):
+        yield from bps.mv(dae.tcb_settings, modified_settings)
+
+        current = yield from bps.rd(dae.tcb_settings)
+        assert current.time_unit == modified_settings.time_unit
+
+    with patch("ibex_bluesky_core.plan_stubs.time_channels_wrapper.ensure_connected"):
+        RE(
+            with_time_channels(
+                _dummy_plan_which_sets_time_units(dae),
+                dae=dae,  # type: ignore
+            )
+        )
+
+    after: DaeTCBSettingsData = RE(bps.rd(dae.tcb_settings)).plan_result  # type: ignore
+    assert after == before
+    assert after.time_unit == expected_time_units
+
+
+def test_dae_table_wrapper(dae: Dae, RE: RunEngine):
+    original_settings = initial_dae_settings
+    modified_settings = DaeSettingsData(
+        wiring_filepath="C:\\somefile.dat",
+        spectra_filepath="C:\\anotherfile.dat",
+        detector_filepath="C:\\anotherfile123.dat",
+    )
+
+    expected_wiring = "NIMROD84modules+9monitors+LAB5Oct2012Wiring.dat"
+    expected_spectra = "NIMROD84modules+9monitors+LAB5Oct2012Spectra.dat"
+    expected_detector = "NIMROD84modules+9monitors+LAB5Oct2012Detector.dat"
+
+    set_mock_value(dae.dae_settings._raw_dae_settings, original_settings)
+
+    before: DaeSettingsData = RE(bps.rd(dae.dae_settings)).plan_result  # type: ignore
+
+    def _dummy_plan_which_sets_wiring_sepectra_detector(dae):
+        yield from bps.mv(dae.dae_settings, modified_settings)
+
+        current = yield from bps.rd(dae.dae_settings)
+        assert current.wiring_filepath == modified_settings.wiring_filepath
+        assert current.spectra_filepath == modified_settings.spectra_filepath
+        assert current.detector_filepath == modified_settings.detector_filepath
+
+    with patch("ibex_bluesky_core.plan_stubs.dae_table_wrapper.ensure_connected"):
+        RE(
+            with_dae_tables(
+                _dummy_plan_which_sets_wiring_sepectra_detector(dae),
+                dae=dae,  # type: ignore
+            )
+        )
+
+    after: DaeSettingsData = RE(bps.rd(dae.dae_settings)).plan_result  # type: ignore
+
+    assert after == before
+    assert after.wiring_filepath.endswith(expected_wiring)  # type: ignore
+    assert after.spectra_filepath.endswith(expected_spectra)  # type: ignore
+    assert after.detector_filepath.endswith(expected_detector)  # type: ignore