doc

doc/plan_stubs/matplotlib_helpers.md

@@ -0,0 +1,51 @@
+# `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:
+
+```
+UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fail.
+  fig, ax = plt.subplots()
+```
+
+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.
+
+```{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.
+```
+
+## `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
+from ibex_bluesky_core.callbacks.plotting import LivePlot
+from bluesky.callbacks import LiveFitPlot
+from bluesky.preprocessors import subs_decorator
+
+def my_plan():
+    # BAD
+    # fig, ax = plt.subplots()
+
+    # GOOD
+    fig, ax = yield from matplotlib_subplots()
+
+    # Pass the matplotlib ax object to other callbacks
+    @subs_decorator([
+        LiveFitPlot(..., ax=ax),
+        LivePlot(..., ax=ax),
+    ])
+    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 new_matplotlib_figure_and_axes
+from ibex_bluesky_core.plan_stubs import matplotlib_subplots
 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 new_matplotlib_figure_and_axes()
+    _, ax = yield from matplotlib_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,28 +1,22 @@
 """Core plan stubs."""
 
-import threading
 from collections.abc import Generator
-from typing import Callable, ParamSpec, TypeVar, cast
+from typing import Any, Callable, ParamSpec, TypeVar, cast
 
 import bluesky.plan_stubs as bps
 import matplotlib.pyplot as plt
-from bluesky.callbacks.mpl_plotting import QtAwareCallback
 from bluesky.utils import Msg
-from event_model import RunStart
-from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 
 P = ParamSpec("P")
 T = TypeVar("T")
 
 
-PYPLOT_SUBPLOTS_TIMEOUT = 5
-
-
 CALL_SYNC_MSG_KEY = "ibex_bluesky_core_call_sync"
+CALL_QT_SAFE_MSG_KEY = "ibex_bluesky_core_call_qt_safe"
 
 
-__all__ = ["call_sync", "new_matplotlib_figure_and_axes"]
+__all__ = ["call_sync", "matplotlib_subplots"]
 
 
 def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Generator[Msg, None, T]:
@@ -56,13 +50,13 @@ def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Genera
     return cast(T, (yield Msg(CALL_SYNC_MSG_KEY, func, *args, **kwargs)))
 
 
-def new_matplotlib_figure_and_axes(
-    *args: list, **kwargs: dict
-) -> Generator[Msg, None, tuple[Figure, Axes]]:
-    """Create a new matplotlib figure and axes, using plt.subplots, from within a plan.
+def matplotlib_subplots(
+    *args: list[Any], **kwargs: dict[Any]
+) -> Generator[Msg, None, tuple[Figure, Any]]:
+    """Create a new matplotlib figure and axes, using matplotlib.pyplot.subplots, from a plan.
 
-    This is done in a Qt-safe way, such that if matplotlib is using the Qt backend then
-    UI operations need to be 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.
 
     Args:
         args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots
@@ -72,35 +66,7 @@ def new_matplotlib_figure_and_axes(
         tuple of (figure, axes) - as per matplotlib.pyplot.subplots()
 
     """
-    yield from bps.null()
-
-    ev = threading.Event()
-    fig: Figure | None = None
-    ax: Axes | None = None
-    ex: BaseException | None = None
-
-    # Slightly hacky, this isn't really a callback per-se but we want to benefit from
-    # bluesky's Qt-matplotlib infrastructure.
-    # This never gets attached to the RunEngine.
-    class _Cb(QtAwareCallback):
-        def start(self, _: RunStart) -> None:
-            nonlocal fig, ax, ex
-            # Note: this is "fast" - so don't need to worry too much about
-            # interruption case.
-            try:
-                fig, ax = plt.subplots(*args, **kwargs)
-            except BaseException as e:
-                ex = e
-            finally:
-                ev.set()
-
-    cb = _Cb()
-    # Send fake event to our callback to trigger it (actual contents unimportant)
-    cb("start", {"time": 0, "uid": ""})
-    if not ev.wait(PYPLOT_SUBPLOTS_TIMEOUT):
-        raise OSError("Could not create matplotlib figure and axes (timeout)")
-    if fig is None or ax is None:
-        raise OSError(
-            "Could not create matplotlib figure and axes (got fig=%s, ax=%s, ex=%s)", fig, ax, ex
-        )
-    return fig, ax
+    yield from bps.clear_checkpoint()
+    return cast(
+        tuple[Figure, Any], (yield Msg(CALL_QT_SAFE_MSG_KEY, plt.subplots, *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_SYNC_MSG_KEY
-from ibex_bluesky_core.run_engine._msg_handlers import call_sync_handler
+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
 
 logger = logging.getLogger(__name__)
 
@@ -97,6 +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.preprocessors.append(functools.partial(bpp.plan_mutator, msg_proc=add_rb_number_processor))
 

src/ibex_bluesky_core/run_engine/_msg_handlers.py

@@ -6,14 +6,19 @@
 import ctypes
 import logging
 import threading
-from asyncio import CancelledError, Event, get_running_loop
+from asyncio import CancelledError, Event, get_running_loop, wait_for
 from typing import Any
 
+from bluesky.callbacks.mpl_plotting import QtAwareCallback
 from bluesky.utils import Msg
+from event_model import RunStart
 
 logger = logging.getLogger(__name__)
 
 
+CALL_QT_SAFE_TIMEOUT = 5
+
+
 class _ExternalFunctionInterrupted(BaseException):
     """An external sync function running in a worker thread is being interrupted."""
 
@@ -94,3 +99,48 @@ def _wrapper() -> Any:  # noqa: ANN401
         logger.debug("Re-raising %s thrown by %s", exc.__class__.__name__, func.__name__)
         raise exc
     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.
+
+    In particular, operations that take longer than CALL_QT_SAFE_TIMEOUT will cause an
+    asyncio TimeoutError (to flag obvious misuse), but won't cancel the underlying task.
+
+    """
+    func = msg.obj
+    done_event = Event()
+    result: Any = None
+    exc: BaseException | None = None
+    loop = get_running_loop()
+
+    # Slightly hacky, this isn't really a callback per-se but we want to benefit from
+    # bluesky's Qt-matplotlib infrastructure.
+    # This never gets attached to the RunEngine.
+    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:
+                result = func(*msg.args, **msg.kwargs)
+            except BaseException as 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)
+    cb("start", {"time": 0, "uid": ""})
+
+    await wait_for(done_event.wait(), CALL_QT_SAFE_TIMEOUT)
+    if exc is not None:
+        raise exc
+    return result