From 2e6cb17e0d7a3cf67f249bc566873c7b51e55c1f Mon Sep 17 00:00:00 2001
From: Tom Willemsen <tom.willemsen@stfc.ac.uk>
Date: Sat, 30 Nov 2024 10:22:06 +0000
Subject: [PATCH] Helper to create a new figure/axes in a qt-safe way

---
 manual_system_tests/dae_scan.py              |  4 +-
 src/ibex_bluesky_core/plan_stubs/__init__.py | 65 +++++++++++++++++++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/manual_system_tests/dae_scan.py b/manual_system_tests/dae_scan.py
index eb9e32a..685a857 100644
--- a/manual_system_tests/dae_scan.py
+++ b/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 new_matplotlib_figure_and_axes
 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 new_matplotlib_figure_and_axes()
+
     lf = LiveFit(
         Linear.fit(), y=reducer.intensity.name, x=block.name, yerr=reducer.intensity_stddev.name
     )
diff --git a/src/ibex_bluesky_core/plan_stubs/__init__.py b/src/ibex_bluesky_core/plan_stubs/__init__.py
index 803e43b..7563020 100644
--- a/src/ibex_bluesky_core/plan_stubs/__init__.py
+++ b/src/ibex_bluesky_core/plan_stubs/__init__.py
@@ -1,17 +1,30 @@
 """Core plan stubs."""
 
-from typing import Callable, Generator, ParamSpec, TypeVar, cast
+import threading
+from collections.abc import Generator
+from typing import 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"
 
 
+__all__ = ["call_sync", "new_matplotlib_figure_and_axes"]
+
+
 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.
 
@@ -41,3 +54,53 @@ def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Genera
     """
     yield from bps.clear_checkpoint()
     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.
+
+    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.
+
+    Args:
+        args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots
+        kwargs: Arbitrary keyword arguments, passed through to matplotlib.pyplot.subplots
+
+    Returns:
+        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