Repository: bluesky/bluesky
Pull request #1639 (2023-12-11 11:07:24)

Add prepare command to RunEngine and bundler, as well as protocol
resolves https://github.com/bluesky/bluesky/issues/1637



Patch:
diff --git a/bluesky/plan_stubs.py b/bluesky/plan_stubs.py
index db1e600e..908dfa4e 100644
--- a/bluesky/plan_stubs.py
+++ b/bluesky/plan_stubs.py
@@ -590,6 +590,39 @@ def input_plan(prompt=''):
     return (yield Msg('input', prompt=prompt))
 
 
+def prepare(obj, *args, group=None, wait=False, **kwargs):
+    """
+    Prepare a device.
+
+    Parameters
+    ----------
+    obj : Preparable
+        Device with 'prepare' method
+    group : string (or any hashable object), optional
+        identifier used by 'wait'
+    wait : boolean, optional
+        If True, wait for completion before processing any more messages.
+        False by default.
+    kwargs
+        passed through to ``obj.prepare()``
+
+    Yields
+    ------
+    msg : Msg
+        Msg('kickoff', obj)
+
+    See Also
+    --------
+    :func:`bluesky.plan_stubs.complete`
+    :func:`bluesky.plan_stubs.collect`
+    :func:`bluesky.plan_stubs.wait`
+    """
+    ret = (yield Msg('prepare', obj, *args, group=group, **kwargs))
+    if wait:
+        yield from _wait(group=group)
+    return ret
+
+
 def kickoff(obj, *, group=None, wait=False, **kwargs):
     """
     Kickoff a fly-scanning device.
diff --git a/bluesky/protocols.py b/bluesky/protocols.py
index 556cc424..412447b3 100644
--- a/bluesky/protocols.py
+++ b/bluesky/protocols.py
@@ -176,6 +176,48 @@ class Triggerable(Protocol):
         ...
 
 
+@runtime_checkable
+class Preparable(Protocol):
+    @abstractmethod
+    def prepare(self, value) -> Status:
+        """Prepare a device for scanning.
+
+        This method provides similar functionality to ``Stageable.stage`` and
+        ``Movable.set``, with key differences:
+
+        ``Stageable.stage``
+        ^^^^^^^^^^^^^^^^^^^
+        Staging a device translates to, "I'm going to use this in a scan, but I'm
+        not sure how". Preparing it translates to, "I'm about to do a step or a fly
+        scan with these parameters". Staging should be universal across many different
+        types of scans, however prepare is specific to an input value passed in.
+
+        ``Movable.set``
+        ^^^^^^^^^^^^^^^
+        For some devices, preparation for a scan could involve multiple soft or
+        hardware signals being configured and/or set. ``prepare`` therefore allows
+        these to be bundled together, along with other logic.
+
+        For example, a Flyable device should have the following methods called on it to
+        perform a fly-scan:
+
+            prepare(flyscan_params)
+            kickoff()
+            complete()
+
+        If the device is a detector, ``collect_asset_docs`` can be called repeatedly
+        while ``complete`` is not done to publish frames. Alternatively, to step-scan a
+        detector,
+
+            prepare(frame_params) to setup N software triggered frames
+            trigger() to take N frames
+            collect_asset_docs() to publish N frames
+
+        Returns a Status that is marked done when the device is ready for a scan.
+        """
+        ...
+
+
 @runtime_checkable
 class Readable(HasName, Protocol):
     @abstractmethod
diff --git a/bluesky/run_engine.py b/bluesky/run_engine.py
index 1ef5f568..c23f3991 100644
--- a/bluesky/run_engine.py
+++ b/bluesky/run_engine.py
@@ -14,7 +14,7 @@ import weakref
 from dataclasses import dataclass
 import typing
 from .bundlers import RunBundler, maybe_await
-from .protocols import (Flyable, Locatable, Movable, Pausable, Readable, Stageable,
+from .protocols import (Flyable, Locatable, Movable, Pausable, Preparable, Readable, Stageable,
                         Status, Stoppable, Triggerable, check_supports)
 
 import concurrent
@@ -477,6 +477,7 @@ class RunEngine:
             'pause': self._pause,
             '_resume_from_suspender': self._resume,
             '_start_suspender': self._start_suspender,
+            'prepare': self._prepare,
             'collect': self._collect,
             'kickoff': self._kickoff,
             'complete': self._complete,
@@ -1987,6 +1988,41 @@ class RunEngine:
             ) from ke
         await current_run.drop(msg)
 
+    async def _prepare(self, msg):
+        """Prepare a flyer for a flyscan
+
+        Expected message object is:
+
+        If `flyer_object` obeys the Preparable protocol, it should have a .prepare
+        method that takes an argument to be set:
+
+            Msg('prepare', flyer_object, value)
+
+        Where value represents an initial state to move the flyer to.
+        """
+        obj = check_supports(msg.obj, Preparable)
+        kwargs = dict(msg.kwargs)
+        group = kwargs.pop('group', None)
+        ret = obj.prepare(*msg.args, **kwargs)
+        p_event = asyncio.Event(**self._loop_for_kwargs)
+        pardon_failures = self._pardon_failures
+
+        def done_callback(status=None):
+            self.log.debug("The object %r reports set is done "
+                           "with status %r", obj, ret.success)
+            self._loop.call_soon_threadsafe(
+                self._status_object_completed, ret, p_event, pardon_failures)
+
+        try:
+            ret.add_callback(done_callback)
+        except AttributeError:
+            # for ophyd < v0.8.0
+            ret.finished_cb = done_callback
+        self._groups[group].add(p_event.wait)
+        self._status_objs[group].add(ret)
+
+        return ret
+
     async def _kickoff(self, msg):
         """Start a flyscan object
 
diff --git a/docs/source/hardware.rst b/docs/source/hardware.rst
index bce58fd9..669fcdde 100644
--- a/docs/source/hardware.rst
+++ b/docs/source/hardware.rst
@@ -255,6 +255,9 @@ the required method.
 .. autoclass:: bluesky.protocols.Hints
     :members:
 
+.. autoclass:: bluesky.protocols.Preparable
+   :members:
+   :show-inheritance:
 
 Checking if an object supports an interface
 -------------------------------------------


Test patch:
diff --git a/bluesky/tests/test_flyer.py b/bluesky/tests/test_flyer.py
index 8c4b5b4f..813b3805 100644
--- a/bluesky/tests/test_flyer.py
+++ b/bluesky/tests/test_flyer.py
@@ -5,6 +5,7 @@ import pytest
 
 from bluesky import Msg
 from bluesky.plans import fly, count
+from bluesky.protocols import Preparable
 from bluesky.run_engine import IllegalMessageSequence
 from bluesky.tests import requires_ophyd
 from bluesky.tests.utils import DocCollector
@@ -418,3 +419,39 @@ def test_describe_config_optional(RE):
 
     assert "simple_flyer" in desc["object_keys"]
     assert "simple_flyer" in desc["configuration"]
+
+
+def test_prepare(RE):
+    """Tests that obj.prepare is correctly called for a fly scan."""
+    class PreparableFlyer(FlyerDevice, Preparable):
+        @call_counter
+        def prepare(self, value):
+            self.value = value
+            return NullStatus()
+
+    flyer = PreparableFlyer(name="flyer")
+
+    # We can prepare the flyer even if no run is open,
+    RE([Msg("prepare", flyer, "abcdefg")])
+
+    assert flyer.value == "abcdefg"
+
+    fly_scan_sequence = [
+        Msg("open_run"),
+        Msg("prepare", flyer, 123),
+        Msg("kickoff", flyer, group="foo"),
+        Msg("wait", group="foo"),
+        Msg("complete", flyer, group="bar"),
+        Msg("wait", group="bar"),
+        Msg("collect", flyer),
+        Msg("close_run"),
+    ]
+
+    RE(fly_scan_sequence)
+
+    assert flyer.value == 123
+
+    assert flyer.call_counts["prepare"] == 2  # called once at start of test
+    assert flyer.call_counts["kickoff"] == 1
+    assert flyer.call_counts["complete"] == 1
+    assert flyer.call_counts["collect"] == 1