ENH: Add block argument to stc.plot() (gh-14105)

doc/changes/dev/14185.newfeature.rst

@@ -0,0 +1 @@
+Add a ``block`` parameter to :func:`mne.viz.plot_source_estimates`, :meth:`mne.SourceEstimate.plot` and :meth:`mne.VolSourceEstimate.plot_3d` to halt execution until the figure is closed, by `Cedric Conday`_.

mne/gui/_coreg.py

@@ -65,7 +65,7 @@
     _plot_mri_fiducials,
     _plot_sensors_3d,
 )
-from ..viz.backends._utils import _qt_app_exec, _qt_safe_window
+from ..viz.backends._utils import _qt_block, _qt_safe_window
 from ..viz.utils import safe_event
 
 
@@ -381,8 +381,8 @@ def _get_default(var, val):
         self._trans_modified = False
         self._mri_fids_modified = False
         self._mri_scale_modified = False
-        if block and self._renderer._kind != "notebook":
-            _qt_app_exec(self._renderer.figure.store["app"])
+        if block:
+            _qt_block(self._renderer)
 
     def _set_subjects_dir(self, subjects_dir):
         if subjects_dir is None or not subjects_dir:

mne/source_estimate.py

@@ -778,6 +778,7 @@ def plot(
         view_layout="vertical",
         add_data_kwargs=None,
         brain_kwargs=None,
+        block=False,
         verbose=None,
     ):
         from .viz import plot_source_estimates
@@ -813,6 +814,7 @@ def plot(
             view_layout=view_layout,
             add_data_kwargs=add_data_kwargs,
             brain_kwargs=brain_kwargs,
+            block=block,
             verbose=verbose,
         )
         return brain
@@ -2345,6 +2347,7 @@ def plot_3d(
         view_layout="vertical",
         add_data_kwargs=None,
         brain_kwargs=None,
+        block=False,
         verbose=None,
     ):
         return super().plot(
@@ -2377,6 +2380,7 @@ def plot_3d(
             view_layout=view_layout,
             add_data_kwargs=add_data_kwargs,
             brain_kwargs=brain_kwargs,
+            block=block,
             verbose=verbose,
         )
 

mne/viz/_3d.py

@@ -2168,6 +2168,7 @@ def _plot_mpl_stc(
     time_viewer=False,
     colorbar=True,
     transparent=True,
+    block=False,
 ):
     """Plot source estimate using mpl."""
     import matplotlib.pyplot as plt
@@ -2307,7 +2308,7 @@ def _plot_mpl_stc(
         cax.tick_params(labelsize=16)
         cb.ax.set_facecolor("0.5")
         cax.set(xlim=(scale_pts[0], scale_pts[2]))
-    plt_show(True)
+    plt_show(True, block=block)
     return fig
 
 
@@ -2416,6 +2417,7 @@ def plot_source_estimates(
     view_layout="vertical",
     add_data_kwargs=None,
     brain_kwargs=None,
+    block=False,
     verbose=None,
 ):
     """Plot SourceEstimate.
@@ -2517,6 +2519,7 @@ def plot_source_estimates(
     %(view_layout)s
     %(add_data_kwargs)s
     %(brain_kwargs)s
+    %(block)s
     %(verbose)s
 
     Returns
@@ -2536,12 +2539,14 @@ def plot_source_estimates(
     - https://openwetware.org/wiki/Beauchamp:FreeSurfer
     """  # noqa: E501
     from ..source_estimate import _BaseSourceEstimate, _check_stc_src
+    from .backends._utils import _qt_block
     from .backends.renderer import _get_3d_backend, use_3d_backend
 
     _check_stc_src(stc, src)
     _validate_type(stc, _BaseSourceEstimate, "stc", "source estimate")
     subjects_dir = get_subjects_dir(subjects_dir=subjects_dir, raise_error=True)
     subject = _check_subject(stc.subject, subject)
+    _validate_type(block, bool, "block")
     _check_option("backend", backend, ["auto", "matplotlib", "pyvistaqt", "notebook"])
     plot_mpl = backend == "matplotlib"
     if not plot_mpl:
@@ -2572,10 +2577,10 @@ def plot_source_estimates(
         transparent=transparent,
     )
     if plot_mpl:
-        return _plot_mpl_stc(stc, spacing=spacing, **kwargs)
+        return _plot_mpl_stc(stc, spacing=spacing, block=block, **kwargs)
     else:
         with use_3d_backend(backend):
-            return _plot_stc(
+            brain = _plot_stc(
                 stc,
                 overlay_alpha=alpha,
                 brain_alpha=alpha,
@@ -2593,6 +2598,9 @@ def plot_source_estimates(
                 title=title,
                 **kwargs,
             )
+        if block:
+            _qt_block(brain._renderer)
+        return brain
 
 
 def _plot_stc(

mne/viz/backends/_utils.py

@@ -275,6 +275,19 @@ def _qt_app_exec(app):
             signal.signal(signal.SIGINT, old_signal)
 
 
+def _qt_block(renderer):
+    """Halt execution until the renderer's window is closed.
+
+    Does nothing for backends that have no Qt application to run, such as the
+    notebook backend.
+    """
+    if renderer._kind == "notebook":
+        return
+    app = renderer.figure.store.get("app")
+    if app is not None:
+        _qt_app_exec(app)
+
+
 def _qt_detect_theme():
     try:
         import darkdetect

mne/viz/backends/tests/test_utils.py

@@ -2,6 +2,8 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
+from types import SimpleNamespace
+
 import numpy as np
 import pytest
 
@@ -11,6 +13,7 @@
     _check_color,
     _get_colormap_from_array,
     _pixmap_to_ndarray,
+    _qt_block,
     _qt_is_dark,
 )
 from mne.viz.utils import _is_dark
@@ -30,6 +33,38 @@ def test_get_colormap_from_array():
     assert isinstance(cmap, ListedColormap)
 
 
+def _fake_renderer(kind, store):
+    return SimpleNamespace(_kind=kind, figure=SimpleNamespace(store=store))
+
+
+def test_qt_block_without_qt_app():
+    """Test that _qt_block is a no-op when there is no Qt app to run."""
+    # the notebook backend never has a Qt application
+    _qt_block(_fake_renderer("notebook", {}))
+    # neither does a renderer whose plotter was supplied by the caller
+    _qt_block(_fake_renderer("qt", {}))
+
+
+def test_qt_block_runs_event_loop():
+    """Test that _qt_block does not return until the Qt application quits."""
+    pytest.importorskip("qtpy")
+    from qtpy.QtCore import QTimer
+    from qtpy.QtWidgets import QApplication
+
+    app = QApplication.instance() or QApplication([])
+    quit_ran = []
+
+    def _quit():
+        quit_ran.append(True)
+        app.quit()
+
+    # If _qt_block returned without running the event loop, the timer would never
+    # fire and quit_ran would still be empty when we check it.
+    QTimer.singleShot(100, _quit)
+    _qt_block(_fake_renderer("qt", {"app": app}))
+    assert quit_ran == [True]
+
+
 def test_check_color():
     """Test color format."""
     assert _check_color("red") == (1.0, 0.0, 0.0)

mne/viz/tests/test_3d.py

@@ -1057,6 +1057,8 @@ def test_process_clim_plot(renderer_interactive, brain_gc):
     brain = stc.plot(**kwargs)
     assert brain.data["center"] is None
     brain.close()
+    with pytest.raises(TypeError, match="block must be an instance of bool"):
+        stc.plot(block="yes", **kwargs)
     brain = stc.plot(clim=dict(pos_lims=(10, 50, 90)), **kwargs)
     assert brain.data["center"] == 0.0
     brain.close()