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

Author: CedricConday

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
 
 
@@ -380,8 +380,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 and self._renderer._kind == "qt":
+            _qt_block(self._renderer.plotter.app_window)
 
     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
@@ -2399,6 +2401,7 @@ def plot_3d(
         view_layout="vertical",
         add_data_kwargs=None,
         brain_kwargs=None,
+        block=False,
         verbose=None,
     ):
         return super().plot(
@@ -2431,6 +2434,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

@@ -2189,6 +2189,7 @@ def _plot_mpl_stc(
     time_viewer=False,
     colorbar=True,
     transparent=True,
+    block=False,
 ):
     """Plot source estimate using mpl."""
     import matplotlib.pyplot as plt
@@ -2328,7 +2329,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
 
 
@@ -2437,6 +2438,7 @@ def plot_source_estimates(
     view_layout="vertical",
     add_data_kwargs=None,
     brain_kwargs=None,
+    block=False,
     verbose=None,
 ):
     """Plot SourceEstimate.
@@ -2538,6 +2540,7 @@ def plot_source_estimates(
     %(view_layout)s
     %(add_data_kwargs)s
     %(brain_kwargs)s
+    %(block)s
     %(verbose)s
 
     Returns
@@ -2557,12 +2560,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:
@@ -2593,10 +2598,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,
@@ -2614,6 +2619,9 @@ def plot_source_estimates(
                 title=title,
                 **kwargs,
             )
+        if block and brain._renderer._kind == "qt":
+            _qt_block(brain.plotter.app_window)
+        return brain
 
 
 def _plot_stc(

mne/viz/tests/test_3d.py

@@ -1149,6 +1149,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()