Helper to create a new figure/axes in a qt-safe way

doc/callbacks/plotting.md

@@ -31,15 +31,23 @@ to plot a scan with a logarithmically-scaled y-axis:
 ```python
 import matplotlib.pyplot as plt
 from ibex_bluesky_core.callbacks.plotting import LivePlot
-# Create a new figure to plot onto.
-plt.figure()
-# Make a new set of axes on that figure
-ax = plt.gca()
-# Set the y-scale to logarithmic
-ax.set_yscale("log")
-# Use the above axes in a LivePlot callback
-plot_callback = LivePlot(y="y_variable", x="x_variable", ax=ax, yerr="yerr_variable")
-# yerr is the uncertanties of each y value, producing error bars
+from ibex_bluesky_core.plan_stubs import call_qt_aware
+
+def plan():
+    # Create a new figure to plot onto.
+    yield from call_qt_aware(plt.figure)
+    # Make a new set of axes on that figure
+    ax = yield from call_qt_aware(plt.gca)
+    # Set the y-scale to logarithmic
+    yield from call_qt_aware(ax.set_yscale, "log")
+    # Use the above axes in a LivePlot callback
+    plot_callback = LivePlot(y="y_variable", x="x_variable", ax=ax, yerr="yerr_variable")
+    # yerr is the uncertanties of each y value, producing error bars
+```
+
+```{note}
+See [docs for `call_qt_aware`](../plan_stubs/matplotlib_helpers.md) for a description of why we need to use 
+`yield from call_qt_aware` rather than calling `matplotlib` functions directly.
 ```
 
 By providing a signal name to the `yerr` argument you can pass uncertainties to LivePlot, by not providing anything for this argument means that no errorbars will be drawn. Errorbars are drawn after each point collected, displaying their standard deviation- uncertainty data is collected from Bluesky event documents and errorbars are updated after every new point added.

doc/plan_stubs/matplotlib_helpers.md

@@ -0,0 +1,50 @@
+# `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 a hang or 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 {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 - for example to create or close a set of axes before passing them to callbacks.
+```
+
+Usage example:
+
+```python
+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 - likely to either crash or hang the plan.
+    # plt.close("all")
+    # fig, ax = plt.subplots()
+
+    # GOOD
+    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([
+        LiveFitPlot(..., ax=ax),
+        LivePlot(..., ax=ax),
+    ])
+    def inner_plan():
+        ...
+```

manual_system_tests/dae_scan.py

@@ -27,6 +27,7 @@
     GoodFramesNormalizer,
 )
 from ibex_bluesky_core.devices.simpledae.waiters import GoodFramesWaiter
+from ibex_bluesky_core.plan_stubs import call_qt_aware
 from ibex_bluesky_core.run_engine import get_run_engine
 
 NUM_POINTS: int = 3
@@ -71,7 +72,8 @@ def dae_scan_plan() -> Generator[Msg, None, None]:
     controller.run_number.set_name("run number")
     reducer.intensity.set_name("normalized counts")
 
-    _, ax = plt.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

@@ -11,6 +11,10 @@
 
 
 CALL_SYNC_MSG_KEY = "ibex_bluesky_core_call_sync"
+CALL_QT_AWARE_MSG_KEY = "ibex_bluesky_core_call_qt_aware"
+
+
+__all__ = ["call_qt_aware", "call_sync"]
 
 
 def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Generator[Msg, None, T]:
@@ -36,9 +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 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.
+
+    Only matplotlib functions may be run using this plan stub.
+
+    Args:
+        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:
+        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(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_SYNC_MSG_KEY
-from ibex_bluesky_core.run_engine._msg_handlers import 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,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_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

@@ -9,7 +9,9 @@
 from asyncio import CancelledError, Event, get_running_loop
 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__)
 
@@ -28,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__)
@@ -94,3 +96,50 @@ 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_aware_handler(msg: Msg) -> Any:  # noqa: ANN401
+    """Handle ibex_bluesky_core.plan_stubs.call_sync."""
+    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
+            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": ""})
+
+    # 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,13 +1,13 @@
 # pyright: reportMissingParameterType=false
-
 import time
 from asyncio import CancelledError
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
+import matplotlib.pyplot as plt
 import pytest
 from bluesky.utils import Msg
 
-from ibex_bluesky_core.plan_stubs import call_sync
+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
 
 
@@ -66,3 +66,56 @@ def f():
     end = time.monotonic()
 
     assert end - start == pytest.approx(1, abs=0.2)
+
+
+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_AWARE_MSG_KEY, f, "foo", keyword_arg="bar"))
+
+    result = RE(plan())
+
+    assert result.plan_result == 123
+
+
+def test_call_qt_aware_throws_exception(RE):
+    def f():
+        raise ValueError("broke it")
+
+    def plan():
+        return (yield Msg(CALL_QT_AWARE_MSG_KEY, f))
+
+    with pytest.raises(ValueError, match="broke it"):
+        RE(plan())
+
+
+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 from call_qt_aware(mock, "all"))
+
+    result = RE(plan())
+    assert result.plan_result == 123
+    mock.assert_called_once_with("all")
+
+
+def test_call_qt_aware_non_matplotlib_function(RE):
+    mock = MagicMock()
+    mock.__module__ = "some_random_module"
+
+    def plan():
+        return (yield from call_qt_aware(mock, "arg", keyword_arg="kwarg"))
+
+    with pytest.raises(
+        ValueError, match="Only matplotlib functions should be passed to call_qt_aware"
+    ):
+        RE(plan())
+
+    mock.assert_not_called()