Repository: bluesky/bluesky
Pull request #1780 (2024-08-05 15:04:11)

Add private attribute for aiding in identifying if a function is a bluesky plan
<!--- Provide a general summary of your changes in the Title above -->

## Description
Adds helper functions that can identify if a function is either a generator co-routine, or has been marked with the `@plan` decorator.

## Motivation and Context
Addresses need for checking whether or not a function is a plan as described in: #1779 


## How Has This Been Tested?
Added a parametrized unit test that passes & checks both positive and negative expected results. All other tests pass without issue (locally, waiting on CI).

```
(venv) [jwlodek@dell-rhel8 bluesky]$ git checkout main
Switched to branch 'main'
Your branch is up to date with 'origin/main'.

(venv) [jwlodek@dell-rhel8 bluesky]$ python
Python 3.11.9 (main, Jun 19 2024, 10:02:06) [GCC 8.5.0 20210514 (Red Hat 8.5.0-22)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import inspect
>>> from bluesky.plan_stubs import mv
>>> inspect.isgeneratorfunction(mv)
False
>>> exit()

(venv) [jwlodek@dell-rhel8 bluesky]$ git checkout custom-wrapper-identify-plan 
Switched to branch 'custom-wrapper-identify-plan'

(venv) [jwlodek@dell-rhel8 bluesky]$ python
Python 3.11.9 (main, Jun 19 2024, 10:02:06) [GCC 8.5.0 20210514 (Red Hat 8.5.0-22)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from bluesky.utils import is_plan
>>> from bluesky.plan_stubs import mv
>>> is_plan(mv)
True
>>> is_plan(print)
False
>>> exit()
```


Patch:
diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py
index fabb780b..c2fc235e 100644
--- a/src/bluesky/utils/__init__.py
+++ b/src/bluesky/utils/__init__.py
@@ -1949,12 +1949,12 @@ class Plan:
         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
@@ -1963,8 +1963,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)


Test patch:
diff --git a/src/bluesky/tests/test_utils.py b/src/bluesky/tests/test_utils.py
index 52d2a81a..c8f4a18c 100644
--- a/src/bluesky/tests/test_utils.py
+++ b/src/bluesky/tests/test_utils.py
@@ -7,11 +7,13 @@ import pytest
 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