Merge pull request bluesky#1780 from jwlodek/custom-wrapper-identify-…

…plan

Add private attribute for aiding in identifying if a function is a bluesky plan

src/bluesky/tests/test_utils.py

@@ -7,11 +7,13 @@
 from cycler import cycler
 
 from bluesky import RunEngine
+from bluesky.plan_stubs import complete_all, mv
 from bluesky.utils import (
     CallbackRegistry,
     Msg,
     ensure_generator,
     is_movable,
+    is_plan,
     merge_cycler,
     plan,
     warn_if_msg_args_or_kwargs,
@@ -536,3 +538,17 @@ def test_warning_behavior(gen_func, iterated):
     else:
         with pytest.warns(RuntimeWarning, match=r".*was never iterated.*"):
             RE(gen_func())
+
+
+@pytest.mark.parametrize(
+    "func, is_plan_result",
+    [
+        (print, False),
+        (mv, True),
+        (complete_all, True),
+        (iterating_plan, True),
+        (sample_plan, True),
+    ],
+)
+def test_check_if_func_is_plan(func, is_plan_result):
+    assert is_plan(func) == is_plan_result

src/bluesky/utils/__init__.py

@@ -1966,12 +1966,12 @@ def throw(self, typ, val=None, tb=None):
         return self._iter.throw(typ, val, tb)
 
 
-def plan(plan):
+def plan(bs_plan):
     """Decorator that warns user if a `yield from` is not called
 
     Parameters
     ----------
-    plan : Generator
+    bs_plan : Generator
         Generator coroutine usually a Bluesky plan
 
     Returns
@@ -1980,8 +1980,27 @@ def plan(plan):
         Wrapped plans
     """
 
-    @wraps(plan)
+    @wraps(bs_plan)
     def wrapper(*args, **kwargs) -> Plan:
-        return Plan(plan, *args, **kwargs)
+        return Plan(bs_plan, *args, **kwargs)
+
+    wrapper._is_plan_ = True
 
     return wrapper
+
+
+def is_plan(bs_plan):
+    """Function that checks if function is a generator, or has been marked as plan
+
+    Parameters
+    ----------
+    bs_plan : Function
+        Typically generator coroutine function / Bluesky plan
+
+    Returns
+    -------
+    boolean
+        True if bs_plan arg is a generator, or the __is_plan__ attribute exists and is True.
+    """
+
+    return inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "_is_plan_", False)