diff --git a/src/ibex_bluesky_core/plan_stubs/__init__.py b/src/ibex_bluesky_core/plan_stubs/__init__.py index 776326f..9d34425 100644 --- a/src/ibex_bluesky_core/plan_stubs/__init__.py +++ b/src/ibex_bluesky_core/plan_stubs/__init__.py @@ -1,6 +1,6 @@ """Core plan stubs.""" -from typing import Callable, Generator, ParamSpec, TypeVar +from typing import Callable, Generator, ParamSpec, TypeVar, cast import bluesky.plan_stubs as bps from bluesky.utils import Msg @@ -12,9 +12,7 @@ CALL_SYNC_MSG_KEY = "ibex_bluesky_core_call_sync" -def call_sync( - func: Callable[P, T], *args: P.args, **kwargs: P.kwargs -) -> Generator[Msg, None, None]: +def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Generator[Msg, None, T]: """Call a synchronous user function in a plan, and returns the result of that call. Attempts to guard against the most common pitfalls of naive implementations, for example: @@ -41,4 +39,4 @@ def call_sync( """ yield from bps.clear_checkpoint() - return (yield Msg(CALL_SYNC_MSG_KEY, func, *args, **kwargs)) + return cast(T, (yield Msg(CALL_SYNC_MSG_KEY, func, *args, **kwargs))) diff --git a/src/ibex_bluesky_core/run_engine/_msg_handlers.py b/src/ibex_bluesky_core/run_engine/_msg_handlers.py index 02c1a1d..d16bcbb 100644 --- a/src/ibex_bluesky_core/run_engine/_msg_handlers.py +++ b/src/ibex_bluesky_core/run_engine/_msg_handlers.py @@ -24,15 +24,13 @@ async def call_sync_handler(msg: Msg) -> Any: # noqa: ANN401 loop = get_running_loop() def _wrapper() -> Any: # noqa: ANN401 + nonlocal ret, exc try: - nonlocal ret ret = func(*msg.args, **msg.kwargs) except _ExternalFunctionInterrupted: pass # Suppress stack traces from our special interruption exception. except BaseException as e: - nonlocal exc exc = e - raise finally: loop.call_soon_threadsafe(done_event.set) @@ -51,7 +49,7 @@ def _wrapper() -> Any: # noqa: ANN401 # be a rather long-running external function, so would prevent the interrupt from working) # - Interrupt but don't actually kill the thread, this leads to a misleading result where # a user gets the shell back but the task is still running in the background - # - Hack around with ctypes to inject an injection into the thread running the external + # - Hack around with ctypes to inject an exception into the thread running the external # function. This is generally frowned upon as a bad idea, but may be the "least bad" # solution to running potentially-blocking user code within a plan. # - Force users to pass in coroutines (which await regularly) or functions which check some