Fix: Bind browser close output to creation context (#14077)

Author: mh105

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`_.

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...")

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)