From 5b50bd78e6e10f6bf4272893c3cbcb05adfa93ac Mon Sep 17 00:00:00 2001 From: Jakub Wlodek Date: Mon, 5 Aug 2024 10:59:15 -0400 Subject: [PATCH 1/4] Adding helper functions for identifying if a function is a bluesky plan --- src/bluesky/tests/test_utils.py | 16 ++++++++++++ src/bluesky/utils/__init__.py | 44 ++++++++++++++++++++++++++++++--- 2 files changed, 56 insertions(+), 4 deletions(-) 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 fabb780b9..4ecfd8c6a 100644 --- a/src/bluesky/utils/__init__.py +++ b/src/bluesky/utils/__init__.py @@ -2,6 +2,7 @@ import asyncio import collections.abc import datetime +import functools import inspect import itertools import operator @@ -1949,12 +1950,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 @@ -1963,8 +1964,43 @@ def plan(plan): Wrapped plans """ - @wraps(plan) + @wraps_plan(bs_plan, plan) def wrapper(*args, **kwargs) -> Plan: - return Plan(plan, *args, **kwargs) + return Plan(bs_plan, *args, **kwargs) return wrapper + + +def mark_as_plan(wrapper, wrapped, decorator, **kwargs): + """Function that adds attr if plan decorator is applied""" + + wrapper = functools.update_wrapper(wrapper, wrapped, **kwargs) + if decorator is plan: + wrapper.__is_plan__ = True + return wrapper + + +def wraps_plan(wrapped, decorator, **kwargs): + """Modified wraps that marks decorated function as bs plan""" + + return functools.partial(mark_as_plan, wrapped=wrapped, decorator=decorator, **kwargs) + + +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. + """ + + if inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "__is_plan__", False): + return True + else: + return False From 2575096b94eae849dbe478f79038874add92e5de Mon Sep 17 00:00:00 2001 From: Jakub Wlodek Date: Tue, 6 Aug 2024 16:04:25 -0400 Subject: [PATCH 2/4] Simplify __is_plan__ tagging logic based on review --- src/bluesky/utils/__init__.py | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py index 4ecfd8c6a..f7420e14e 100644 --- a/src/bluesky/utils/__init__.py +++ b/src/bluesky/utils/__init__.py @@ -2,7 +2,6 @@ import asyncio import collections.abc import datetime -import functools import inspect import itertools import operator @@ -1964,28 +1963,15 @@ def plan(bs_plan): Wrapped plans """ - @wraps_plan(bs_plan, plan) + @wraps(bs_plan) def wrapper(*args, **kwargs) -> Plan: return Plan(bs_plan, *args, **kwargs) - return wrapper - - -def mark_as_plan(wrapper, wrapped, decorator, **kwargs): - """Function that adds attr if plan decorator is applied""" + wrapper.__is_plan__ = True - wrapper = functools.update_wrapper(wrapper, wrapped, **kwargs) - if decorator is plan: - wrapper.__is_plan__ = True return wrapper -def wraps_plan(wrapped, decorator, **kwargs): - """Modified wraps that marks decorated function as bs plan""" - - return functools.partial(mark_as_plan, wrapped=wrapped, decorator=decorator, **kwargs) - - def is_plan(bs_plan): """Function that checks if function is a generator, or has been marked as plan From 56a3531f512c396e055059f5a97785c2d0ba393b Mon Sep 17 00:00:00 2001 From: Jakub Wlodek Date: Wed, 4 Sep 2024 14:00:16 -0400 Subject: [PATCH 3/4] Update to use single underscore --- src/bluesky/utils/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py index f7420e14e..2c811754d 100644 --- a/src/bluesky/utils/__init__.py +++ b/src/bluesky/utils/__init__.py @@ -1967,7 +1967,7 @@ def plan(bs_plan): def wrapper(*args, **kwargs) -> Plan: return Plan(bs_plan, *args, **kwargs) - wrapper.__is_plan__ = True + wrapper._is_plan_ = True return wrapper @@ -1986,7 +1986,7 @@ def is_plan(bs_plan): True if bs_plan arg is a generator, or the __is_plan__ attribute exists and is True. """ - if inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "__is_plan__", False): + if inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "_is_plan_", False): return True else: return False From fd06b5de55b24855b7323d828bef7df5757ba940 Mon Sep 17 00:00:00 2001 From: Jakub Wlodek Date: Wed, 4 Sep 2024 21:19:29 -0400 Subject: [PATCH 4/4] Simplify return condition Co-authored-by: Thomas A Caswell --- src/bluesky/utils/__init__.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py index 2c811754d..c2fc235e6 100644 --- a/src/bluesky/utils/__init__.py +++ b/src/bluesky/utils/__init__.py @@ -1986,7 +1986,4 @@ def is_plan(bs_plan): True if bs_plan arg is a generator, or the __is_plan__ attribute exists and is True. """ - if inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "_is_plan_", False): - return True - else: - return False + return inspect.isgeneratorfunction(bs_plan) or getattr(bs_plan, "_is_plan_", False)