Merge pull request #17802 from richardsheridan/tk_backend_figuremanager_quit

QuLogic · web-flow · commit 97fb36168958 · 2020-08-06T21:14:29.000-04:00
fix FigureManagerTk close behavior if embedded in Tk App
diff --git a/lib/matplotlib/backends/_backend_tk.py b/lib/matplotlib/backends/_backend_tk.py
@@ -398,6 +398,8 @@ class FigureManagerTk(FigureManagerBase):
         The tk.Window
     """
 
+    _owns_mainloop = False
+
     def __init__(self, canvas, num, window):
         FigureManagerBase.__init__(self, canvas, num)
         self.window = window
@@ -442,9 +444,8 @@ def show(self):
         with _restore_foreground_window_at_end():
             if not self._shown:
                 def destroy(*args):
-                    self.window = None
                     Gcf.destroy(self)
-                self.canvas._tkcanvas.bind("<Destroy>", destroy)
+                self.window.protocol("WM_DELETE_WINDOW", destroy)
                 self.window.deiconify()
             else:
                 self.canvas.draw_idle()
@@ -454,15 +455,13 @@ def destroy(*args):
             self._shown = True
 
     def destroy(self, *args):
-        if self.window is not None:
-            #self.toolbar.destroy()
-            if self.canvas._idle_callback:
-                self.canvas._tkcanvas.after_cancel(self.canvas._idle_callback)
-            self.window.destroy()
-        if Gcf.get_num_fig_managers() == 0:
-            if self.window is not None:
-                self.window.quit()
-        self.window = None
+        if self.canvas._idle_callback:
+            self.canvas._tkcanvas.after_cancel(self.canvas._idle_callback)
+
+        self.window.destroy()
+
+        if self._owns_mainloop and not Gcf.get_num_fig_managers():
+            self.window.quit()
 
     def get_window_title(self):
         return self.window.wm_title()
@@ -890,4 +889,12 @@ def trigger_manager_draw(manager):
     def mainloop():
         managers = Gcf.get_all_fig_managers()
         if managers:
-            managers[0].window.mainloop()
+            first_manager = managers[0]
+            manager_class = type(first_manager)
+            if manager_class._owns_mainloop:
+                return
+            manager_class._owns_mainloop = True
+            try:
+                first_manager.window.mainloop()
+            finally:
+                manager_class._owns_mainloop = False
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -64,9 +64,10 @@ def _get_running_interactive_framework():
         return "wx"
     tkinter = sys.modules.get("tkinter")
     if tkinter:
+        codes = {tkinter.mainloop.__code__, tkinter.Misc.mainloop.__code__}
         for frame in sys._current_frames().values():
             while frame:
-                if frame.f_code == tkinter.mainloop.__code__:
+                if frame.f_code in codes:
                     return "tk"
                 frame = frame.f_back
     if 'matplotlib.backends._macosx' in sys.modules:
diff --git a/lib/matplotlib/tests/test_backend_tk.py b/lib/matplotlib/tests/test_backend_tk.py
@@ -1,5 +1,11 @@
-import pytest
+import os
+import subprocess
+import sys
+import tkinter
+
 import numpy as np
+import pytest
+
 from matplotlib import pyplot as plt
 
 
@@ -26,3 +32,71 @@ def evil_blit(photoimage, aggimage, offsets, bboxptr):
                       np.ones((4, 4, 4)),
                       (0, 1, 2, 3),
                       bad_boxes)
+
+
+@pytest.mark.backend('TkAgg', skip_on_importerror=True)
+def test_figuremanager_preserves_host_mainloop():
+    success = False
+
+    def do_plot():
+        plt.figure()
+        plt.plot([1, 2], [3, 5])
+        plt.close()
+        root.after(0, legitimate_quit)
+
+    def legitimate_quit():
+        root.quit()
+        nonlocal success
+        success = True
+
+    root = tkinter.Tk()
+    root.after(0, do_plot)
+    root.mainloop()
+
+    assert success
+
+
+@pytest.mark.backend('TkAgg', skip_on_importerror=True)
+@pytest.mark.flaky(reruns=3)
+def test_figuremanager_cleans_own_mainloop():
+    script = '''
+import tkinter
+import time
+import matplotlib.pyplot as plt
+import threading
+from matplotlib.cbook import _get_running_interactive_framework
+
+root = tkinter.Tk()
+plt.plot([1, 2, 3], [1, 2, 5])
+
+def target():
+    while not 'tk' == _get_running_interactive_framework():
+        time.sleep(.01)
+    plt.close()
+    if show_finished_event.wait():
+        print('success')
+
+show_finished_event = threading.Event()
+thread = threading.Thread(target=target, daemon=True)
+thread.start()
+plt.show(block=True)  # testing if this function hangs
+show_finished_event.set()
+thread.join()
+
+'''
+    try:
+        proc = subprocess.run(
+            [sys.executable, "-c", script],
+            env={**os.environ,
+                 "MPLBACKEND": "TkAgg",
+                 "SOURCE_DATE_EPOCH": "0"},
+            timeout=10,
+            stdout=subprocess.PIPE,
+            universal_newlines=True,
+            check=True
+        )
+    except subprocess.TimeoutExpired:
+        pytest.fail("Most likely plot.show(block=True) hung")
+    except subprocess.CalledProcessError:
+        pytest.fail("Subprocess failed to test intended behavior")
+    assert proc.stdout.count("success") == 1
