Repository: bluesky/bluesky (python)

Issue:
Let stage and unstage return Status
At the moment `stage()` and `unstage()` are sync and return a `List[Any]` of the devices that have been staged.

Discussions with @tacaswell and @danielballan suggest this return value is rarely used, and it would be more valuable to stage multiple devices at the same time.

Suggest we change the protocols to be:
- `def stage() -> Status`
- `def unstage() -> Status`

Need to change https://github.com/bluesky/bluesky/blob/d2b6bddd4415cdb2946fe8d552bd704bc9932c11/bluesky/run_engine.py#L2347-L2378 to look at the return values. If it is a `Status` object then track and return it like in https://github.com/bluesky/bluesky/blob/d2b6bddd4415cdb2946fe8d552bd704bc9932c11/bluesky/run_engine.py#L2130-L2160 
but if not then just return the result for backwards compatibility.

Also need to change https://github.com/bluesky/bluesky/blob/d2b6bddd4415cdb2946fe8d552bd704bc9932c11/bluesky/plan_stubs.py#L692-L729 to also call `wait()` if the result of the call is a Status message like we do in https://github.com/bluesky/bluesky/blob/d2b6bddd4415cdb2946fe8d552bd704bc9932c11/bluesky/plan_stubs.py#L913-L918



Pull request description:
ENH: make stage and unstage return status objects
Let stage and unstage return Status objects, so multiple objects can be staged at once.

## Description

I have made changes to  `_stage` and `_unstage` to allow Stageables to return status objects from their stage and unstage methods. The changes mimic `_trigger`. 

I have changed the type hints on the Stageable protocol.

I have created a test for staging 2 Stageables at once.  

I have also changed the plan stubs accordingly.

## Motivation and Context

resolves #1559 

This change will allow the run engine to stage multiple devices without blocking. And these can then be awaited but a wait message.

## How Has This Been Tested?

These changes have been tested by a new test `test_stage_and_unstage_status_objects()`.
Which stages/unstages 2 dummy Stageables then awaits their completion.

The new `_stage` and `_unstage` methods use `isinstance()` to check to see if the output of an objects stage/unstage method is a status object. If it is not they will return the result preserving previous behaviour for Stageables that returns lists of staged objects.


Patch:
diff --git a/bluesky/plan_stubs.py b/bluesky/plan_stubs.py
index 16dfbc02..ceb3511d 100644
--- a/bluesky/plan_stubs.py
+++ b/bluesky/plan_stubs.py
@@ -13,7 +13,7 @@ try:
 except ImportError:
     from toolz import partition
 
-from .protocols import Locatable, Triggerable
+from .protocols import Locatable, Triggerable, Status
 from .utils import (
     get_hinted_fields,
     merge_cycler,
@@ -689,44 +689,142 @@ def configure(obj, *args, **kwargs):
     return (yield Msg('configure', obj, *args, **kwargs))
 
 
-def stage(obj):
+def stage(obj, *, group=None, wait=None):
     """
     'Stage' a device (i.e., prepare it for use, 'arm' it).
 
     Parameters
     ----------
     obj : Device
+    group : string (or any hashable object), optional
+        identifier used by 'wait'; None by default
+    wait : boolean, optional
+        If True, wait for completion before processing any more messages.
+        False by default.
 
     Yields
     ------
     msg : Msg
-        Msg('stage', obj)
 
     See Also
     --------
     :func:`bluesky.plan_stubs.unstage`
+    :func:`bluesky.plan_stubs.stage_all`
+    """
+    ret = yield Msg('stage', obj, group=group)
+    old_style = not isinstance(ret, Status)
+    if old_style:
+        if (wait is None) or wait:
+            # Old-style devices will just block. We do not need to explicitly wait.
+            pass
+        else:  # wait is False-y
+            # No way to tell old-style devices not to wait
+            raise RuntimeError(f"{obj}: Is an old style device and cannot be told not to wait")
+    else:
+        if wait:
+            yield Msg('wait', None, group=group)
+    return ret
+
+
+def stage_all(*args, group=None):
+    """
+    'Stage' one or more devices (i.e., prepare them for use, 'arm' them).
+
+    Parameters
+    ----------
+    args :
+        device1, device2, device3, ...
+    group : string (or any hashable object), optional
+        identifier used by 'wait'; None by default
+
+    Yields
+    ------
+    msg : Msg
+
+    See Also
+    --------
+    :func:`bluesky.plan_stubs.stage`
+    :func:`bluesky.plan_stubs.unstage_all`
     """
-    return (yield Msg('stage', obj))
+    group = group or str(uuid.uuid4())
+    status_objects = []
 
+    for obj in args:
+        ret = yield Msg('stage', obj, group=group)
+        if isinstance(ret, Status):
+            status_objects.append(ret)
 
-def unstage(obj):
+    if status_objects:
+        yield Msg('wait', None, group=group)
+
+
+def unstage(obj, *, group=None, wait=None):
     """
     'Unstage' a device (i.e., put it in standby, 'disarm' it).
 
     Parameters
     ----------
     obj : Device
+    group : string (or any hashable object), optional
+        identifier used by 'wait'; None by default
+    wait : boolean, optional
+        If True, wait for completion before processing any more messages.
+        False by default.
 
     Yields
     ------
     msg : Msg
-        Msg('unstage', obj)
 
     See Also
     --------
     :func:`bluesky.plan_stubs.stage`
+    :func:`bluesky.plan_stubs.unstage_all`
+    """
+    ret = yield Msg('unstage', obj, group=group)
+    old_style = not isinstance(ret, Status)
+    if old_style:
+        if (wait is None) or wait:
+            # Old-style devices will just block. We do not need to explicitly wait.
+            pass
+        else:
+            # No way to tell old-style devices not to wait
+            raise RuntimeError(f"{obj}: Is an old style device and cannot be told not to wait")
+    else:
+        if wait:
+            yield Msg('wait', None, group=group)
+    return ret
+
+
+def unstage_all(*args, group=None):
     """
-    return (yield Msg('unstage', obj))
+    'Unstage' one or more devices (i.e., put them in standby, 'disarm' them).
+
+    Parameters
+    ----------
+    args :
+        device1, device2, device3, ...
+    group : string (or any hashable object), optional
+        identifier used by 'wait'; None by default
+
+    Yields
+    ------
+    msg : Msg
+
+    See Also
+    --------
+    :func:`bluesky.plan_stubs.unstage`
+    :func:`bluesky.plan_stubs.stage_all`
+    """
+    group = group or str(uuid.uuid4())
+    status_objects = []
+
+    for obj in args:
+        ret = yield Msg('unstage', obj, group=group)
+        if isinstance(ret, Status):
+            status_objects.append(ret)
+
+    if status_objects:
+        yield Msg('wait', None, group=group)
 
 
 def subscribe(name, func):
diff --git a/bluesky/protocols.py b/bluesky/protocols.py
index b605c54a..82ce604d 100644
--- a/bluesky/protocols.py
+++ b/bluesky/protocols.py
@@ -341,22 +341,20 @@ class Stageable(Protocol):
     # TODO: we were going to extend these to be able to return plans, what
     # signature should they have?
     @abstractmethod
-    def stage(self) -> List[Any]:
+    def stage(self) -> Union[Status, List[Any]]:
         """An optional hook for "setting up" the device for acquisition.
 
-        It should return a list of devices including itself and any other
-        devices that are staged as a result of staging this one.
-        (The ``parent`` attribute expresses this relationship: a device should
-        be staged/unstaged whenever its parent is staged/unstaged.)
+        It should return a ``Status`` that is marked done when the device is
+        done staging.
         """
         ...
 
     @abstractmethod
-    def unstage(self) -> List[Any]:
+    def unstage(self) -> Union[Status, List[Any]]:
         """A hook for "cleaning up" the device after acquisition.
 
-        It should return a list of devices including itself and any other
-        devices that are unstaged as a result of unstaging this one.
+        It should return a ``Status`` that is marked done when the device is finished
+        unstaging.
         """
         ...
 
diff --git a/bluesky/run_engine.py b/bluesky/run_engine.py
index 173c0006..59095432 100644
--- a/bluesky/run_engine.py
+++ b/bluesky/run_engine.py
@@ -15,7 +15,7 @@ from dataclasses import dataclass
 import typing
 from .bundlers import RunBundler, maybe_await
 from .protocols import (Flyable, Locatable, Movable, Pausable, Readable, Stageable,
-                        Stoppable, Triggerable, check_supports)
+                        Status, Stoppable, Triggerable, check_supports)
 
 import concurrent
 
@@ -135,6 +135,7 @@ class LoggingPropertyMachine(PropertyMachine):
     """expects object to have a `log` attribute
     and a `state_hook` attribute that is ``None`` or a callable with signature
     ``f(value, old_value)``"""
+
     def __init__(self, machine_type):
         super().__init__(machine_type)
 
@@ -1133,7 +1134,6 @@ class RunEngine:
             explanation of why the suspension has been requested
 
         """
-
         print("Suspending....To get prompt hit Ctrl-C twice to pause.")
         ts = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
         print("Suspension occurred at {}.".format(ts))
@@ -2140,22 +2140,8 @@ class RunEngine:
         group = kwargs.pop('group', None)
         warn_if_msg_args_or_kwargs(msg, obj.trigger, msg.args, kwargs)
         ret = obj.trigger(*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 trigger 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)
+        self._add_status_to_group(obj=obj, status_object=ret, group=group, action="trigger")
 
         return ret
 
@@ -2344,6 +2330,25 @@ class RunEngine:
             await current_run.configure(msg)
         return old, new
 
+    def _add_status_to_group(self, obj: typing.Any, status_object: Status,
+                             group: str, action: str) -> None:
+        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 %r is "
+                           "done with status %r.", obj, action, status_object.success)
+            self._loop.call_soon_threadsafe(
+                self._status_object_completed, status_object, p_event, pardon_failures)
+
+        try:
+            status_object.add_callback(done_callback)
+        except AttributeError:
+            # for ophyd < v0.8.0
+            status_object.finished_cb = done_callback
+        self._groups[group].add(p_event.wait)
+        self._status_objs[group].add(status_object)
+
     async def _stage(self, msg):
         """Instruct the RunEngine to stage the object
 
@@ -2355,10 +2360,17 @@ class RunEngine:
         # If an object has no 'stage' method, assume there is nothing to do.
         if not isinstance(obj, Stageable):
             return []
-        result = obj.stage()
+        group = kwargs.pop('group', None)
+        ret = obj.stage()
         self._staged.add(obj)  # add first in case of failure below
         await self._reset_checkpoint_state_coro()
-        return result
+
+        if not isinstance(ret, Status):
+            return ret
+
+        self._add_status_to_group(obj=obj, status_object=ret, group=group, action="stage")
+
+        return ret
 
     async def _unstage(self, msg):
         """Instruct the RunEngine to unstage the object
@@ -2371,11 +2383,18 @@ class RunEngine:
         # If an object has no 'unstage' method, assume there is nothing to do.
         if not isinstance(obj, Stageable):
             return []
-        result = obj.unstage()
+        group = kwargs.pop('group', None)
+        ret = obj.unstage()
         # use `discard()` to ignore objects that are not in the staged set.
         self._staged.discard(obj)
         await self._reset_checkpoint_state_coro()
-        return result
+
+        if not isinstance(ret, Status):
+            return ret
+
+        self._add_status_to_group(obj=obj, status_object=ret, group=group, action="unstage")
+
+        return ret
 
     async def _stop(self, msg):
         """


Test patch:
diff --git a/bluesky/tests/test_new_examples.py b/bluesky/tests/test_new_examples.py
index e8c4b519..e6959de9 100644
--- a/bluesky/tests/test_new_examples.py
+++ b/bluesky/tests/test_new_examples.py
@@ -1,5 +1,6 @@
 import asyncio
 from collections import defaultdict
+import time as ttime
 from typing import Dict
 import pytest
 from bluesky import Msg, RunEngineInterrupted
@@ -24,7 +25,9 @@ from bluesky.plan_stubs import (
     collect,
     configure,
     stage,
+    stage_all,
     unstage,
+    unstage_all,
     subscribe,
     unsubscribe,
     install_suspender,
@@ -106,8 +109,10 @@ import threading
      (collect, ('foo',), {}, [Msg('collect', 'foo',
                                   stream=False, return_payload=True)]),
      (configure, ('det', 1), {'a': 2}, [Msg('configure', 'det', 1, a=2)]),
-     (stage, ('det',), {}, [Msg('stage', 'det')]),
-     (unstage, ('det',), {}, [Msg('unstage', 'det')]),
+     (stage, ('det',), {}, [Msg('stage', 'det', group=None)]),
+     (stage, ('det',), {"group": "A"}, [Msg('stage', 'det', group="A")]),
+     (unstage, ('det',), {}, [Msg('unstage', 'det', group=None)]),
+     (unstage, ('det',), {"group": "A"}, [Msg('unstage', 'det', group="A")]),
      (subscribe, ('all', 'func_placeholder'), {}, [Msg('subscribe', None,
                                                        'func_placeholder',
                                                        'all')]),
@@ -807,3 +812,91 @@ def test_per_step(RE, hw):
     with pytest.raises(TypeError) as excinfo:
         RE(scan([hw.det], hw.motor, -1, 1, 3, per_step=bad_sig))
     assert excinfo.match("per_step must be a callable with the signature")
+
+
+def test_stage_all_and_unstage_all(RE):
+
+    from ophyd import StatusBase
+
+    staged = {}
+    unstaged = {}
+
+    # Test support for old and new style devices
+    class OldStyleDummy:
+        def __init__(self, name: str) -> None:
+            self.name: str = name
+
+        def stage(self):
+            staged[self.name] = True
+            ttime.sleep(0.5)
+            return [self]
+
+        def unstage(self):
+            unstaged[self.name] = True
+            ttime.sleep(0.5)
+            return [self]
+
+    class NewStyleDummy:
+        def __init__(self, name: str) -> None:
+            self.name: str = name
+
+        def _callback(self, d: dict, st: Status):
+            d[self.name] = True
+            st.set_finished()
+
+        def stage(self) -> Status:
+            st = StatusBase()
+            threading.Timer(0.5, self._callback, args=[staged, st]).start()
+            return st
+
+        def unstage(self) -> Status:
+            st = StatusBase()
+            threading.Timer(0.5, self._callback, args=[unstaged, st]).start()
+            return st
+
+    olddummy1 = OldStyleDummy("o1")
+    olddummy2 = OldStyleDummy("o2")
+    newdummy1 = NewStyleDummy("n1")
+    newdummy2 = NewStyleDummy("n2")
+
+    def plan():
+        yield from stage_all(olddummy1, olddummy2, newdummy1, newdummy2)
+        assert list(staged.keys()) == ["o1", "o2", "n1", "n2"]
+
+        yield from unstage_all(olddummy1, olddummy2, newdummy1, newdummy2)
+        assert list(unstaged.keys()) == ["o1", "o2", "n1", "n2"]
+
+    start = ttime.monotonic()
+    RE(plan())
+    stop = ttime.monotonic()
+
+    assert 3 < stop - start < 4
+
+
+def test_old_style_wait(RE):
+
+    class OldStyleDummy:
+        def __init__(self, name: str) -> None:
+            self.name: str = name
+
+        def stage(self):
+            ttime.sleep(0.5)
+            return [self]
+
+        def unstage(self):
+            ttime.sleep(0.5)
+            return [self]
+
+    dummy = OldStyleDummy("o1")
+
+    def stage_plan():
+        yield from stage(dummy, wait=False)
+
+    def unstage_plan():
+        yield from unstage(dummy, wait=False)
+
+    with pytest.raises(RuntimeError):
+        RE(stage_plan())
+
+    with pytest.raises(RuntimeError):
+        RE(unstage_plan())
diff --git a/bluesky/tests/test_run_engine.py b/bluesky/tests/test_run_engine.py
index fba4cefe..22f6bc1f 100644
--- a/bluesky/tests/test_run_engine.py
+++ b/bluesky/tests/test_run_engine.py
@@ -8,6 +8,7 @@ import sys
 from collections import defaultdict
 import time as ttime
 import pytest
+from bluesky.protocols import Status
 from bluesky.tests import requires_ophyd
 from bluesky.run_engine import (RunEngineStateMachine,
                                 TransitionError, IllegalMessageSequence,
@@ -386,6 +387,64 @@ def test_pause_resume_devices(RE):
     assert 'dummy' in resumed
 
 
+def test_stage_and_unstage_status_objects(RE):
+
+    from ophyd import StatusBase
+
+    staged = {}
+    unstaged = {}
+    started_stage = {}
+    started_unstage = {}
+
+    class Dummy:
+        def __init__(self, name: str) -> None:
+            self.name: str = name
+
+        def _callback(self, d: dict, st: Status):
+            d[self.name] = True
+            st.set_finished()
+
+        def stage(self) -> Status:
+            st = StatusBase()
+            started_stage[self.name] = True
+            threading.Timer(1, self._callback, args=[staged, st]).start()
+            return st
+
+        def unstage(self) -> Status:
+            st = StatusBase()
+            started_unstage[self.name] = True
+            threading.Timer(1, self._callback, args=[unstaged, st]).start()
+            return st
+
+    dummy1 = Dummy("d1")
+    dummy2 = Dummy("d2")
+
+    def my_plan():
+        yield Msg("stage", dummy1, group="stage_group")
+        yield Msg("stage", dummy2, group="stage_group")
+
+        assert "d1" in started_stage and "d2" in started_stage
+        assert not staged
+
+        yield Msg("wait", group="stage_group")
+        assert "d1" in staged and "d2" in staged
+
+        yield Msg("unstage", dummy1, group="unstage_group")
+        yield Msg("unstage", dummy2, group="unstage_group")
+
+        assert "d1" in started_unstage and "d2" in started_unstage
+        assert not unstaged
+
+        yield Msg("wait", group="unstage_group")
+        assert "d1" in unstaged and "d2" in unstaged
+
+    start = ttime.monotonic()
+    RE(my_plan())
+    stop = ttime.monotonic()
+
+    assert 2 < stop - start < 3
+
+
 def test_bad_call_args(RE):
     with pytest.raises(Exception):
         RE(53)