Repository: bluesky/bluesky
Pull request #1663 (2024-03-07 15:43:11)

Alter ko and complete plan stubs to take multiple flyables
This adds the multiple object functionality to the kickoff and complete plan stubs. This will neaten up the plans.


danielballan: I have been pondering the trade offs of making `msg.obj` a list (in some cases) or packing additional objects into `msg.args`. I think this approach is the less messy option of the two.
abbiemery: Changes made, but I feel like we should leave this one until the branch is merged and is in main. There is a log going in which needs a review, a lot of test cases and people who really need #1651 in, so we can make an alpha-release and can therefore make a proper ophyd-async release.

There is no real harm in leaving this, other than it makes the plans slightly messier. Something we can put in a following release.
coretl: I think we are clear that we want to planstubs to kickoff (and complete) multiple devices at once. The questions are:
- Should `bps.kickoff` to work with both a single and multiple devices?
- If so, should the single case move from returning a single Status to returning a tuple of Statuses?
- If not, should we make a `bps.kickoff_all` like [`bps.stage_all`](https://github.com/bluesky/bluesky/blob/f03451a7e188df867df7fc86c70e3efe65a02709/src/bluesky/plan_stubs.py#L787)

@tacaswell @danielballan I feel like we had the "should `func(obj, *other_objs)` always return a tuple of return values, even if there are no `other_objs`?" conversation before, but I can't remember where...
coretl: Decided on a `kickoff_all` and `complete_all` and then we can use different default args so it waits by default

Patch:
diff --git a/src/bluesky/plan_stubs.py b/src/bluesky/plan_stubs.py
index d60bc032..d087151c 100644
--- a/src/bluesky/plan_stubs.py
+++ b/src/bluesky/plan_stubs.py
@@ -5,6 +5,7 @@ import uuid
 import warnings
 from collections.abc import Iterable
 from functools import reduce
+from typing import List
 
 from cycler import cycler
 
@@ -14,7 +15,7 @@ try:
 except ImportError:
     from toolz import partition
 
-from .protocols import Locatable, Status, Triggerable
+from .protocols import Flyable, Locatable, Status, Triggerable, check_supports
 from .utils import (
     Msg,
     all_safe_rewind,
@@ -624,12 +625,11 @@ def prepare(obj, *args, group=None, wait=False, **kwargs):
 
 def kickoff(obj, *, group=None, wait=False, **kwargs):
     """
-    Kickoff a fly-scanning device.
+    Kickoff one fly-scanning device.
 
     Parameters
     ----------
-    obj : fly-able
-        Device with 'kickoff', and 'complete' methods.
+    obj : fly-able Device with 'kickoff', and 'complete' methods.
     group : string (or any hashable object), optional
         identifier used by 'wait'.
     wait : boolean, optional
@@ -655,11 +655,51 @@ def kickoff(obj, *, group=None, wait=False, **kwargs):
     return ret
 
 
+def kickoff_all(*args, group=None, wait=True, **kwargs):
+    """
+    Kickoff one or more fly-scanning devices.
+
+    Parameters
+    ----------
+    *args : Any fly-able
+        Device with 'kickoff', and 'complete' methods.
+    group : string (or any hashable object), optional
+        identifier used by 'wait'.
+    wait : boolean, optional
+        If True, wait for completion before processing any more messages.
+        True by default.
+    kwargs
+        passed through to 'kickoff' for each device
+
+    Yields
+    ------
+    msg : Msg
+        Msg('kickoff', obj)
+
+    See Also
+    --------
+    :func:`bluesky.plan_stubs.complete`
+    :func:`bluesky.plan_stubs.collect`
+    :func:`bluesky.plan_stubs.wait`
+    """
+    objs = [check_supports(arg, Flyable) for arg in args]
+    group = group or str(uuid.uuid4())
+    statuses: List[Status] = []
+
+    for obj in objs:
+        ret = yield Msg("kickoff", obj, group=group, **kwargs)
+        statuses.append(ret)
+    if wait:
+        yield from _wait(group=group)
+
+    return tuple(statuses)
+
+
 def complete(obj, *, group=None, wait=False, **kwargs):
     """
-    Tell a flyer, 'stop collecting, whenever you are ready'.
+    Tell a flyable, 'stop collecting, whenever you are ready'.
 
-    The flyer returns a status object. Some flyers respond to this
+    A flyable returns a status object. Some flyers respond to this
     command by stopping collection and returning a finished status
     object immediately. Other flyers finish their given course and
     finish whenever they finish, irrespective of when this command is
@@ -694,6 +734,52 @@ def complete(obj, *, group=None, wait=False, **kwargs):
     return ret
 
 
+def complete_all(*args, group=None, wait=False, **kwargs):
+    """
+    Tell one or more flyable objects, 'stop collecting, whenever you are ready'.
+
+    A flyable returns a status object. Some flyers respond to this
+    command by stopping collection and returning a finished status
+    object immediately. Other flyers finish their given course and
+    finish whenever they finish, irrespective of when this command is
+    issued.
+
+    Parameters
+    ----------
+    *args : Any fly-able
+        Device with 'kickoff' and 'complete' methods.
+    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 'complete' for each device
+
+    Yields
+    ------
+    msg : Msg
+        a 'complete' Msg and maybe a 'wait' message
+
+    See Also
+    --------
+    :func:`bluesky.plan_stubs.kickoff`
+    :func:`bluesky.plan_stubs.collect`
+    :func:`bluesky.plan_stubs.wait`
+    """
+    objs = [check_supports(arg, Flyable) for arg in args]
+    group = group or str(uuid.uuid4())
+    statuses: List[Status] = []
+
+    for obj in objs:
+        ret = yield Msg("complete", obj, group=group, **kwargs)
+        statuses.append(ret)
+    if wait:
+        yield from _wait(group=group)
+
+    return tuple(statuses)
+
+
 def collect(obj, *args, stream=False, return_payload=True, name=None):
     """
     Collect data cached by one or more fly-scanning devices and emit documents.


Test patch:
diff --git a/src/bluesky/tests/test_flyer.py b/src/bluesky/tests/test_flyer.py
index 79d5956e..4ff17a80 100644
--- a/src/bluesky/tests/test_flyer.py
+++ b/src/bluesky/tests/test_flyer.py
@@ -7,6 +7,7 @@ from ophyd import Device
 from ophyd.sim import NullStatus, TrivialFlyer
 
 from bluesky import Msg
+from bluesky.plan_stubs import close_run, complete, complete_all, kickoff, kickoff_all, open_run, wait
 from bluesky.plans import count, fly
 from bluesky.protocols import Preparable
 from bluesky.run_engine import IllegalMessageSequence
@@ -456,3 +457,82 @@ def test_prepare(RE):
     assert flyer.call_counts["kickoff"] == 1
     assert flyer.call_counts["complete"] == 1
     assert flyer.call_counts["collect"] == 1
+
+
+def test_kickoff_complete_all_single_device(RE):
+    flyer1 = FlyerDevice(name="flyer1")
+
+    def fly_plan():
+        yield from open_run()
+        yield from kickoff_all(flyer1, group="foo")
+        yield from wait(group="foo")
+        yield from complete_all(flyer1, group="bar")
+        yield from wait(group="bar")
+        yield from close_run()
+
+    RE(fly_plan())
+
+    assert flyer1.call_counts["kickoff"] == 1
+    assert flyer1.call_counts["complete"] == 1
+
+
+def test_kickoff_complete_all_take_multiple_devices(RE):
+    flyer1 = FlyerDevice(name="flyer1")
+    flyer2 = FlyerDevice(name="flyer2")
+
+    def fly_plan():
+        yield from open_run()
+        yield from kickoff_all(flyer1, flyer2, group="foo")
+        yield from wait(group="foo")
+        yield from complete_all(flyer1, flyer2, group="bar")
+        yield from wait(group="bar")
+        yield from close_run()
+
+    RE(fly_plan())
+
+    assert flyer1.call_counts["kickoff"] == 1
+    assert flyer2.call_counts["kickoff"] == 1
+    assert flyer1.call_counts["complete"] == 1
+    assert flyer2.call_counts["complete"] == 1
+
+
+def test_kickoff_all_complete_seperately(RE):
+    flyer1 = FlyerDevice(name="flyer1")
+    flyer2 = FlyerDevice(name="flyer2")
+
+    def fly_plan():
+        yield from open_run()
+        yield from kickoff_all(flyer1, flyer2, group="foo")
+        yield from wait(group="foo")
+        yield from complete(flyer1, group="bar")
+        yield from complete(flyer2, group="bar")
+        yield from wait(group="bar")
+        yield from close_run()
+
+    RE(fly_plan())
+
+    assert flyer1.call_counts["kickoff"] == 1
+    assert flyer2.call_counts["kickoff"] == 1
+    assert flyer1.call_counts["complete"] == 1
+    assert flyer2.call_counts["complete"] == 1
+
+
+def test_kickoff_seperately_complete_all(RE):
+    flyer1 = FlyerDevice(name="flyer1")
+    flyer2 = FlyerDevice(name="flyer2")
+
+    def fly_plan():
+        yield from open_run()
+        yield from kickoff(flyer1, group="foo")
+        yield from kickoff(flyer2, group="foo")
+        yield from wait(group="foo")
+        yield from complete_all(flyer1, flyer2, group="bar")
+        yield from wait(group="bar")
+        yield from close_run()
+
+    RE(fly_plan())
+
+    assert flyer1.call_counts["kickoff"] == 1
+    assert flyer2.call_counts["kickoff"] == 1
+    assert flyer1.call_counts["complete"] == 1
+    assert flyer2.call_counts["complete"] == 1