diff --git a/doc/changes/dev/14077.bugfix.rst b/doc/changes/dev/14077.bugfix.rst new file mode 100644 index 00000000000..48b2b92c55e --- /dev/null +++ b/doc/changes/dev/14077.bugfix.rst @@ -0,0 +1 @@ +Fix Qt browser close messages from :meth:`mne.io.Raw.plot` appearing in the wrong Jupyter notebook cell, by `Mingjian He`_. diff --git a/mne/viz/_figure.py b/mne/viz/_figure.py index faea296ff88..f1f17898fa7 100644 --- a/mne/viz/_figure.py +++ b/mne/viz/_figure.py @@ -10,6 +10,7 @@ from abc import ABC, abstractmethod from collections import OrderedDict from contextlib import contextmanager +from contextvars import copy_context from copy import deepcopy from itertools import cycle @@ -56,6 +57,7 @@ def __init__(self, **kwargs): from ..preprocessing import ICA self.backend_name = None + self._close_context = copy_context() self._data = None self._times = None @@ -490,6 +492,13 @@ def _redraw(self, update_data=True, annotations=False): def _close(self, event=None): """Handle close events (via keypress or window [x]).""" + # As specified by PEP 567: https://peps.python.org/pep-0567/#asyncio + # we explicitly retain the python Context used to create the figure + # in order to route stdout on close within IPykernel. See gh #14077 + self._close_context.run(self._close_impl, event) + + def _close_impl(self, event=None): + """Handle close events in the context that created the browser.""" from matplotlib.pyplot import close logger.debug(f"Closing {self.mne.instance_type} browser...") diff --git a/mne/viz/tests/test_figure.py b/mne/viz/tests/test_figure.py index 936c4cab0e9..df5588a80cc 100644 --- a/mne/viz/tests/test_figure.py +++ b/mne/viz/tests/test_figure.py @@ -2,12 +2,15 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. +from contextvars import ContextVar + +import matplotlib.pyplot as plt import numpy as np import pytest from mne import create_info from mne.io import RawArray -from mne.viz._figure import _get_browser +from mne.viz._figure import _get_browser, use_browser_backend def test_browse_figure_constructor(): @@ -23,3 +26,28 @@ def test_browse_figure_requires_two_timepoints(): assert len(raw.times) == 1 with pytest.raises(ValueError, match="at least two time points"): _get_browser(show=False, block=False, inst=raw) + + +def test_browse_figure_close_context(): + """Test that deferred browser close uses its creation context.""" + marker = ContextVar("browser_context", default=None) + created_token = marker.set("created") + try: + info = create_info(ch_names=["CH1"], sfreq=100.0, ch_types="eeg") + raw = RawArray(np.zeros((1, 100)), info) + with use_browser_backend("matplotlib"): + fig = raw.plot(show=False) + + current_token = marker.set("current") + observed = [] + close_impl = fig._close_impl + fig._close_impl = lambda event=None: observed.append(marker.get()) + try: + fig._close() + finally: + fig._close_impl = close_impl + plt.close(fig) + marker.reset(current_token) + assert observed == ["created"] + finally: + marker.reset(created_token)