Simplify/genericise

doc/plan_stubs/matplotlib_helpers.md

@@ -1,7 +1,7 @@
-# `matplotlib` helpers
+# `call_qt_aware` (matplotlib helpers)
 
 When attempting to use `matplotlib` UI functions directly in a plan, and running `matplotlib` using a `Qt`
-backend (e.g. in a standalone shell outside IBEX), you may see an error of the form:
+backend (e.g. in a standalone shell outside IBEX), you may see a hang or an error of the form:
 
 ```
 UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fail.
@@ -11,34 +11,34 @@ UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fa
 This is because the `RunEngine` runs plans in a worker thread, not in the main thread, which then requires special
 handling when calling functions that will update a UI.
 
-The following plan stubs provide Qt-safe wrappers around some matplotlib functions to avoid this error.
+The {py:obj}`ibex_bluesky_core.plan_stubs.call_qt_aware` plan stub can call `matplotlib` functions in a
+Qt-aware context, which allows them to be run directly from a plan. It allows the same arguments and 
+keyword-arguments as the underlying matplotlib function it is passed.
 
 ```{note}
 Callbacks such as `LivePlot` and `LiveFitPlot` already route UI calls to the appropriate UI thread by default. 
 The following plan stubs are only necessary if you need to call functions which will create or update a matplotlib
-plot from a plan directly.
+plot from a plan directly - for example to create or close a set of axes before passing them to callbacks.
 ```
 
-## `matplotlib_subplots`
-
-The {py:obj}`ibex_bluesky_core.plan_stubs.matplotlib_subplots` plan stub is a Qt-safe wrapper 
-around `matplotlib.pyplot.subplots()`. It allows the same arguments and keyword-arguments as the 
-underlying matplotlib function.
-
 Usage example:
 
 ```python
-from ibex_bluesky_core.plan_stubs import matplotlib_subplots
+import matplotlib.pyplot as plt
+from ibex_bluesky_core.plan_stubs import call_qt_aware
 from ibex_bluesky_core.callbacks.plotting import LivePlot
 from bluesky.callbacks import LiveFitPlot
 from bluesky.preprocessors import subs_decorator
 
+
 def my_plan():
-    # BAD
+    # BAD - likely to either crash or hang the plan.
+    # plt.close("all")
     # fig, ax = plt.subplots()
 
     # GOOD
-    fig, ax = yield from matplotlib_subplots()
+    yield from call_qt_aware(plt.close, "all")
+    fig, ax = yield from call_qt_aware(plt.subplots)
 
     # Pass the matplotlib ax object to other callbacks
     @subs_decorator([
@@ -48,4 +48,3 @@ def my_plan():
     def inner_plan():
         ...
 ```
-

manual_system_tests/dae_scan.py

@@ -27,7 +27,7 @@
     GoodFramesNormalizer,
 )
 from ibex_bluesky_core.devices.simpledae.waiters import GoodFramesWaiter
-from ibex_bluesky_core.plan_stubs import matplotlib_subplots
+from ibex_bluesky_core.plan_stubs import call_qt_aware
 from ibex_bluesky_core.run_engine import get_run_engine
 
 NUM_POINTS: int = 3
@@ -72,7 +72,7 @@ def dae_scan_plan() -> Generator[Msg, None, None]:
     controller.run_number.set_name("run number")
     reducer.intensity.set_name("normalized counts")
 
-    _, ax = yield from matplotlib_subplots()
+    _, ax = yield from call_qt_aware(plt.subplots)
 
     lf = LiveFit(
         Linear.fit(), y=reducer.intensity.name, x=block.name, yerr=reducer.intensity_stddev.name

src/ibex_bluesky_core/plan_stubs/__init__.py

@@ -1,22 +1,20 @@
 """Core plan stubs."""
 
 from collections.abc import Generator
-from typing import Any, Callable, ParamSpec, TypeVar, cast
+from typing import Callable, ParamSpec, TypeVar, cast
 
 import bluesky.plan_stubs as bps
-import matplotlib.pyplot as plt
 from bluesky.utils import Msg
-from matplotlib.figure import Figure
 
 P = ParamSpec("P")
 T = TypeVar("T")
 
 
 CALL_SYNC_MSG_KEY = "ibex_bluesky_core_call_sync"
-CALL_QT_SAFE_MSG_KEY = "ibex_bluesky_core_call_qt_safe"
+CALL_QT_AWARE_MSG_KEY = "ibex_bluesky_core_call_qt_aware"
 
 
-__all__ = ["call_sync", "matplotlib_subplots"]
+__all__ = ["call_qt_aware", "call_sync"]
 
 
 def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Generator[Msg, None, T]:
@@ -42,32 +40,41 @@ def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Genera
 
     Args:
         func: A callable to run.
-        args: Arbitrary arguments to be passed to the wrapped function
-        kwargs: Arbitrary keyword arguments to be passed to the wrapped function
+        *args: Arbitrary arguments to be passed to the wrapped function
+        **kwargs: Arbitrary keyword arguments to be passed to the wrapped function
+
+    Returns:
+        The return value of the wrapped function
 
     """
     yield from bps.clear_checkpoint()
     return cast(T, (yield Msg(CALL_SYNC_MSG_KEY, func, *args, **kwargs)))
 
 
-def matplotlib_subplots(
-    *args: Any,  # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API.
-    **kwargs: Any,  # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API.
-) -> Generator[Msg, None, tuple[Figure, Any]]:
-    """Create a new matplotlib figure and axes, using matplotlib.pyplot.subplots, from a plan.
+def call_qt_aware(
+    func: Callable[P, T], *args: P.args, **kwargs: P.kwargs
+) -> Generator[Msg, None, T]:
+    """Call a matplotlib function in a Qt-aware context, from within a plan.
+
+    If matplotlib is using a Qt backend then UI operations are run on the Qt thread via Qt signals.
 
-    This is done in a Qt-safe way, such that if matplotlib is using a Qt backend then
-    UI operations are run on the Qt thread via Qt signals.
+    Only matplotlib functions may be run using this plan stub.
 
     Args:
-        args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots
-        kwargs: Arbitrary keyword arguments, passed through to matplotlib.pyplot.subplots
+        func: A matplotlib function reference.
+        *args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots
+        **kwargs: Arbitrary keyword arguments, passed through to matplotlib.pyplot.subplots
+
+    Raises:
+        ValueError: if the passed function is not a matplotlib function.
 
     Returns:
-        tuple of (figure, axes) - as per matplotlib.pyplot.subplots()
+        The return value of the wrapped function
 
     """
+    # Limit potential for misuse - constrain to just running matplotlib functions.
+    if not getattr(func, "__module__", "").startswith("matplotlib"):
+        raise ValueError("Only matplotlib functions should be passed to call_qt_aware")
+
     yield from bps.clear_checkpoint()
-    return cast(
-        tuple[Figure, Any], (yield Msg(CALL_QT_SAFE_MSG_KEY, plt.subplots, *args, **kwargs))
-    )
+    return cast(T, (yield Msg(CALL_QT_AWARE_MSG_KEY, func, *args, **kwargs)))

src/ibex_bluesky_core/run_engine/__init__.py

@@ -17,8 +17,8 @@
 __all__ = ["get_run_engine"]
 
 
-from ibex_bluesky_core.plan_stubs import CALL_QT_SAFE_MSG_KEY, CALL_SYNC_MSG_KEY
-from ibex_bluesky_core.run_engine._msg_handlers import call_qt_safe_handler, call_sync_handler
+from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, CALL_SYNC_MSG_KEY
+from ibex_bluesky_core.run_engine._msg_handlers import call_qt_aware_handler, call_sync_handler
 
 logger = logging.getLogger(__name__)
 
@@ -97,7 +97,7 @@ def get_run_engine() -> RunEngine:
     RE.subscribe(log_callback)
 
     RE.register_command(CALL_SYNC_MSG_KEY, call_sync_handler)
-    RE.register_command(CALL_QT_SAFE_MSG_KEY, call_qt_safe_handler)
+    RE.register_command(CALL_QT_AWARE_MSG_KEY, call_qt_aware_handler)
 
     RE.preprocessors.append(functools.partial(bpp.plan_mutator, msg_proc=add_rb_number_processor))
 

src/ibex_bluesky_core/run_engine/_msg_handlers.py

@@ -3,11 +3,10 @@
 Not intended for user use.
 """
 
-import asyncio
 import ctypes
 import logging
 import threading
-from asyncio import CancelledError, Event, get_running_loop, wait_for
+from asyncio import CancelledError, Event, get_running_loop
 from typing import Any
 
 from bluesky.callbacks.mpl_plotting import QtAwareCallback
@@ -17,9 +16,6 @@
 logger = logging.getLogger(__name__)
 
 
-CALL_QT_SAFE_TIMEOUT = 5
-
-
 class _ExternalFunctionInterrupted(BaseException):
     """An external sync function running in a worker thread is being interrupted."""
 
@@ -34,7 +30,7 @@ async def call_sync_handler(msg: Msg) -> Any:  # noqa: ANN401
 
     def _wrapper() -> Any:  # noqa: ANN401
         nonlocal ret, exc
-        logger.info("Running '{func.__name__}' with args=({msg.args}), kwargs=({msg.kwargs})")
+        logger.info("Running '%s' with args=(%s), kwargs=(%s)", func.__name__, msg.args, msg.kwargs)
         try:
             ret = func(*msg.args, **msg.kwargs)
             logger.debug("Running '%s' successful", func.__name__)
@@ -102,18 +98,8 @@ def _wrapper() -> Any:  # noqa: ANN401
     return ret
 
 
-async def call_qt_safe_handler(msg: Msg) -> Any:  # noqa: ANN401
-    """Handle ibex_bluesky_core.plan_stubs.call_qt_safe.
-
-    This functionality does not get exposed in a generic way to users, as this is
-    a tricky area. This *relies* on the passed function being "fast", but
-    we would have no way to prove that for a generic user-supplied function. So we only
-    expose known cases, like matplotlib.pyplot.subplots for example.
-
-    Slow functions will cause various undesirable side effects, like stalling the event loop
-    and interfering with ctrl-c handling.
-
-    """
+async def call_qt_aware_handler(msg: Msg) -> Any:  # noqa: ANN401
+    """Handle ibex_bluesky_core.plan_stubs.call_sync."""
     func = msg.obj
     done_event = Event()
     result: Any = None
@@ -126,28 +112,34 @@ async def call_qt_safe_handler(msg: Msg) -> Any:  # noqa: ANN401
     class _Cb(QtAwareCallback):
         def start(self, doc: RunStart) -> None:
             nonlocal result, exc
-            # Note: Qt/UI operations must be "fast", so don't worry too much about timeout or
-            # interruption cases here.
-            # Any attempt to forcibly interrupt a function while it's doing UI operations/using
-            # Qt signals is highly likely to be a bad idea. Don't do that here.
             try:
+                logger.info(
+                    "Running '%s' with args=(%s), kwargs=(%s) (Qt)",
+                    func.__name__,
+                    msg.args,
+                    msg.kwargs,
+                )
                 result = func(*msg.args, **msg.kwargs)
+                logger.debug("Running '%s' (Qt) successful", func.__name__)
             except BaseException as e:
+                logger.error(
+                    "Running '%s' failed with %s: %s", func.__name__, e.__class__.__name__, e
+                )
                 exc = e
             finally:
                 loop.call_soon_threadsafe(done_event.set)
 
     cb = _Cb()
     # Send fake event to our callback to trigger it (actual contents unimportant)
+    # If not using Qt, this will run synchronously i.e. block until complete
+    # If using Qt, this will be sent off to the Qt teleporter which will execute it asynchronously,
+    # and we have to wait for the event to be set.
     cb("start", {"time": 0, "uid": ""})
 
-    try:
-        await wait_for(done_event.wait(), CALL_QT_SAFE_TIMEOUT)
-    except asyncio.TimeoutError as e:  # pragma: no cover (CI on linux doesn't have a UI/Qt backend)
-        raise TimeoutError(
-            f"Long-running function '{func.__name__}' passed to call_qt_safe_handler. Functions "
-            f"passed to call_qt_safe_handler must be faster than {CALL_QT_SAFE_TIMEOUT}s."
-        ) from e
+    # Attempting to forcibly interrupt a function while it's doing UI operations/using
+    # Qt signals is highly likely to be a bad idea. Don't do that here. No special ctrl-c handling.
+    await done_event.wait()
+
     if exc is not None:
         raise exc
     return result

tests/test_plan_stubs.py

@@ -1,14 +1,13 @@
 # pyright: reportMissingParameterType=false
-import re
 import time
 from asyncio import CancelledError
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
-import matplotlib
+import matplotlib.pyplot as plt
 import pytest
 from bluesky.utils import Msg
 
-from ibex_bluesky_core.plan_stubs import CALL_QT_SAFE_MSG_KEY, call_sync, matplotlib_subplots
+from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, call_qt_aware, call_sync
 from ibex_bluesky_core.run_engine._msg_handlers import call_sync_handler
 
 
@@ -69,54 +68,54 @@ def f():
     assert end - start == pytest.approx(1, abs=0.2)
 
 
-def test_call_qt_safe_returns_result(RE):
+def test_call_qt_aware_returns_result(RE):
     def f(arg, keyword_arg):
         assert arg == "foo"
         assert keyword_arg == "bar"
         return 123
 
     def plan():
-        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f, "foo", keyword_arg="bar"))
+        return (yield Msg(CALL_QT_AWARE_MSG_KEY, f, "foo", keyword_arg="bar"))
 
     result = RE(plan())
 
     assert result.plan_result == 123
 
 
-def test_call_qt_safe_throws_exception(RE):
+def test_call_qt_aware_throws_exception(RE):
     def f():
         raise ValueError("broke it")
 
     def plan():
-        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f))
+        return (yield Msg(CALL_QT_AWARE_MSG_KEY, f))
 
     with pytest.raises(ValueError, match="broke it"):
         RE(plan())
 
 
-@pytest.mark.skipif("qt" not in matplotlib.get_backend().lower(), reason="Qt not available")
-def test_call_qt_safe_blocking_causes_descriptive_timeout_error(RE):
-    def f():
-        time.sleep(0.2)
+def test_call_qt_aware_matplotlib_function(RE):
+    mock = MagicMock(spec=plt.close)
+    mock.__module__ = "matplotlib.pyplot"
+    mock.return_value = 123
 
     def plan():
-        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f))
+        return (yield from call_qt_aware(mock, "all"))
 
-    with patch("ibex_bluesky_core.run_engine._msg_handlers.CALL_QT_SAFE_TIMEOUT", new=0.1):
-        with pytest.raises(
-            TimeoutError,
-            match=re.escape(
-                "Long-running function 'f' passed to call_qt_safe_handler. "
-                "Functions passed to call_qt_safe_handler must be faster than 0.1s."
-            ),
-        ):
-            RE(plan())
+    result = RE(plan())
+    assert result.plan_result == 123
+    mock.assert_called_once_with("all")
 
 
-def test_matplotlib_subplots_calls_pyplot_subplots(RE):
+def test_call_qt_aware_non_matplotlib_function(RE):
+    mock = MagicMock()
+    mock.__module__ = "some_random_module"
+
     def plan():
-        return (yield from matplotlib_subplots("foo", keyword="bar"))
+        return (yield from call_qt_aware(mock, "arg", keyword_arg="kwarg"))
 
-    with patch("ibex_bluesky_core.plan_stubs.plt.subplots") as mock_pyplot_subplots:
+    with pytest.raises(
+        ValueError, match="Only matplotlib functions should be passed to call_qt_aware"
+    ):
         RE(plan())
-        mock_pyplot_subplots.assert_called_once_with("foo", keyword="bar")
+
+    mock.assert_not_called()