Repository: bluesky/bluesky
Pull request #1879 (2025-02-03 18:13:59)

FIX: Plan wrapper to deal with closing from stop/abort
<!--- Provide a general summary of your changes in the Title above -->

## Description
<!--- Describe your changes in detail -->
Similar to the explicit inclusion of `send` and `throw` in the `Plan` wrapper, `close` was added as a generator method. This allows pchain to make use of `close` during `RunEngineControlException` raised by `stop` and `abort`. 

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->
Stopping a plan built with `bpp.pchain` would raise warnings that the plan didn't use yield from. E.g.,
```
RE.stop()
...
  File "<ipython-input-11-5cf04c04aa5b>", line 1, in <module>
    Eref = 933; RE(bpp.pchain(asc([sclr], pgm.energy, Eref-5, Eref+15 , 41, md={'purpose':'Energy check at Cu L3 edge on Cu metal'}), bps.sleep(5), mv(pgm.energy, Eref)))
RuntimeWarning: plan `sleep` was never iterated, did you mean to use `yield from`?
  warnings.warn(warning_message, RuntimeWarning, stacklevel=1)
```
## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->
A test was added specifically for the stop behavior. 

<!--
## Screenshots (if appropriate):
-->



Patch:
diff --git a/src/bluesky/preprocessors.py b/src/bluesky/preprocessors.py
index 4f46ef72..f949d246 100644
--- a/src/bluesky/preprocessors.py
+++ b/src/bluesky/preprocessors.py
@@ -5,7 +5,16 @@ from functools import wraps
 
 from bluesky.protocols import Locatable
 
-from .plan_stubs import close_run, declare_stream, mv, open_run, pause, stage_all, trigger_and_read, unstage_all
+from .plan_stubs import (
+    close_run,
+    declare_stream,
+    mv,
+    open_run,
+    pause,
+    stage_all,
+    trigger_and_read,
+    unstage_all,
+)
 from .utils import (
     Msg,
     RunEngineControlException,
@@ -18,9 +27,7 @@ from .utils import (
     separate_devices,
     single_gen,
 )
-from .utils import (
-    short_uid as _short_uid,
-)
+from .utils import short_uid as _short_uid
 
 
 def plan_mutator(plan, msg_proc):
@@ -280,7 +287,7 @@ def pchain(*args):
     """Like `itertools.chain` but using `yield from`
 
     This ensures than `.send` works as expected and the underlying
-    plans get the return values
+    plans get the return values. Provides cleanup for Interruptions.
 
     Parameters
     ----------
@@ -293,8 +300,13 @@ def pchain(*args):
         The messages from each plan in turn
     """
     rets = deque()
-    for p in args:
-        rets.append((yield from p))
+    try:
+        for p in args:
+            rets.append((yield from p))
+    except RunEngineControlException as e:
+        for p in args:
+            p.close()
+        raise e
     return tuple(rets)
 
 
@@ -557,7 +569,13 @@ def finalize_wrapper(plan, final_plan, *, pause_for_debug=False):
 
 
 def contingency_wrapper(
-    plan, *, except_plan=None, else_plan=None, final_plan=None, pause_for_debug=False, auto_raise=True
+    plan,
+    *,
+    except_plan=None,
+    else_plan=None,
+    final_plan=None,
+    pause_for_debug=False,
+    auto_raise=True,
 ):
     """try...except...else...finally helper
 
@@ -1120,7 +1138,10 @@ def relative_set_wrapper(plan, devices=None):
         seen = msg.obj in initial_positions
         if (msg.command == "set") and eligible and not seen:
             return (
-                pchain(__read_and_stash_a_motor(msg.obj, initial_positions, coupled_parents), single_gen(msg)),
+                pchain(
+                    __read_and_stash_a_motor(msg.obj, initial_positions, coupled_parents),
+                    single_gen(msg),
+                ),
                 None,
             )
         else:
@@ -1158,7 +1179,10 @@ def reset_positions_wrapper(plan, devices=None):
         seen = msg.obj in initial_positions
         if (msg.command == "set") and eligible and not seen:
             return (
-                pchain(__read_and_stash_a_motor(msg.obj, initial_positions, coupled_parents), single_gen(msg)),
+                pchain(
+                    __read_and_stash_a_motor(msg.obj, initial_positions, coupled_parents),
+                    single_gen(msg),
+                ),
                 None,
             )
         else:
diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py
index 52cd278b..6a3fdbbb 100644
--- a/src/bluesky/utils/__init__.py
+++ b/src/bluesky/utils/__init__.py
@@ -1968,6 +1968,10 @@ class Plan:
         self._stack = None
         return self._iter.throw(typ, val, tb)
 
+    def close(self):
+        self._stack = None
+        return self._iter.close()
+
 
 def plan(bs_plan):
     """Decorator that warns user if a `yield from` is not called


Test patch:
diff --git a/src/bluesky/tests/test_utils.py b/src/bluesky/tests/test_utils.py
index 3db74438..4745dcdd 100644
--- a/src/bluesky/tests/test_utils.py
+++ b/src/bluesky/tests/test_utils.py
@@ -9,8 +9,9 @@ import numpy as np
 import pytest
 from cycler import cycler
 
-from bluesky import RunEngine
+from bluesky import RunEngine, RunEngineInterrupted
 from bluesky.plan_stubs import complete_all, mv
+from bluesky.preprocessors import pchain
 from bluesky.run_engine import WaitForTimeoutError
 from bluesky.utils import (
     AsyncInput,
@@ -545,6 +546,36 @@ def test_warning_behavior(gen_func, iterated):
             RE(gen_func())
 
 
+def pause_plan():
+    return (yield Msg("pause"))
+
+
+def pchain_plan():
+    yield from pchain(sample_plan(), pause_plan(), sample_plan())
+
+
+def test_warnings_with_interruption():
+    RE = RunEngine()
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        with pytest.raises(RunEngineInterrupted):
+            RE(pchain_plan())
+        assert RE.state == "paused"
+        RE.stop()
+        assert not record, "There should be no warnings if properly closed"
+
+
+def test_plan_wrapper_close():
+    plan = sample_plan()
+    plan.close()
+    with pytest.raises(StopIteration):
+        next(plan)
+    plan = pchain_plan()
+    plan.close()
+    with pytest.raises(StopIteration):
+        next(plan)
+
+
 @pytest.mark.parametrize(
     "func, is_plan_result",
     [
@@ -570,7 +601,10 @@ def test_async_input_does_not_block_event_loop(RE, capsys):
         # Patch stdin.readline as a *blocking* function that takes a while to return - the wait_for
         # should time out before this returns, so the input_completed_event should never have a
         # chance to be set.
-        with patch("bluesky.utils.sys.stdin.readline", side_effect=lambda: time.sleep(3 * a_short_time)):
+        with patch(
+            "bluesky.utils.sys.stdin.readline",
+            side_effect=lambda: time.sleep(3 * a_short_time),
+        ):
             await ai("prompt: ")
 
         input_completed_event.set()