diff --git a/src/bluesky/tests/test_utils.py b/src/bluesky/tests/test_utils.py index 52d2a81a4..c8f4a18ce 100644 --- a/src/bluesky/tests/test_utils.py +++ b/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 diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py index c33a00205..c7c839b60 100644 --- a/src/bluesky/utils/__init__.py +++ b/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)