Merge pull request #267 from ISISComputingGroup/ticket164wrapper

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

Author: Tom-Willemsen

doc/plan_stubs/plan_wrappers.md

@@ -0,0 +1,87 @@
+# Plan wrappers
+
+Plan wrappers that temporarily modify [DAE (Data Acquisition Electronics)](/devices/dae.md) 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
+
+A function that wraps a plan to temporarily modify the DAE table.
+
+API reference: {py:obj}`dae_table_wrapper<ibex_bluesky_core.plan_stubs.with_dae_tables>`
+
+```python
+RE(
+    with_dae_tables(
+        bps.null(),
+        dae=dae,
+        new_settings=modified_settings
+    )
+)
+```
+```python
+def plan():
+    yield from with_dae_tables(scan(...), dae=dae, new_settings=modified_settings)
+```
+
+(where `modified_settings` is a dataset in the form {py:obj}`DaeSettingsData <ibex_bluesky_core.devices.dae.DaeSettingsData>`)
+
+
+
+### Num Periods
+
+A function that wraps a plan to temporarily modify the number of periods.
+
+API reference: {py:obj}`num_periods_wrapper<ibex_bluesky_core.plan_stubs.with_num_periods>`
+
+```python
+RE(
+    with_num_periods(
+        bps.null(),
+        dae=dae,
+        number_of_periods=1000 
+    )
+)
+```
+```python
+def plan():
+    yield from with_num_periods(scan(...), dae=dae, number_of_periods=1000)
+```
+
+
+### Time Channels
+A function that wraps a plan to temporarily modify the time channels boundaries.
+
+API reference: {py:obj}`time_channels_wrapper<ibex_bluesky_core.plan_stubs.with_time_channels>`
+
+```python
+RE(
+    with_time_channels(
+        bps.null(),
+        dae=dae,
+        new_settings=modified_settings
+    )
+)
+```
+```python
+def plan():
+    yield from with_time_channels(scan(...), dae=dae, new_settings=modified_settings)
+```
+(where `modified_settings` is a dataset in the form {py:obj}`DaeTCBSettingsData<ibex_bluesky_core.devices.dae.DaeTCBSettingsData>`)
+
+## Usage
+
+To use these wrappers, the plan written by the user must be wrapped by the function within the RunEngine:
+
+``` python
+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
+
+def plan():
+    yield from with_num_periods(scan(...), dae=dae, number_of_periods=1000)
+```
+
+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 …ibex_bluesky_core/plan_stubs/__init__.pysrc/ibex_bluesky_core/plan_stubs.py renamed to src/ibex_bluesky_core/plan_stubs/__init__.py src/ibex_bluesky_core/plan_stubs.py renamed to 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__)
@@ -31,6 +34,9 @@
     "prompt_user_for_choice",
     "redefine_motor",
     "redefine_refl_parameter",
+    "with_dae_tables",
+    "with_num_periods",
+    "with_time_channels",
 ]
 
 

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)
+
+        return (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)
+
+        return (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)
+
+        return (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)