Make image_comparison work even without the autoclose fixture.

anntzer · anntzer · commit b35f0bcf9bbf · 2021-10-01T20:09:03.000+02:00
... by capturing only the images created during the execution of the
test, and closing them after the comparison.
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -14,10 +14,7 @@
 import matplotlib.style
 import matplotlib.units
 import matplotlib.testing
-from matplotlib import cbook
-from matplotlib import ft2font
-from matplotlib import pyplot as plt
-from matplotlib import ticker
+from matplotlib import cbook, ft2font, pyplot as plt, ticker, _pylab_helpers
 from .compare import comparable_formats, compare_images, make_test_filename
 from .exceptions import ImageComparisonFailure
 
@@ -129,6 +126,28 @@ def remove_ticks(ax):
         remove_ticks(ax)
 
 
+@contextlib.contextmanager
+def _collect_new_figures():
+    """
+    After::
+
+        with _collect_new_figures() as figs: some_code()
+
+    *figs* contains the figures that have been created during the execution of
+    ``some_code``, sorted by figure number.
+    """
+    managers = _pylab_helpers.Gcf.figs
+    preexisting = [manager for manager in managers.values()]
+    new_figs = []
+    try:
+        yield new_figs
+    finally:
+        new_managers = sorted([manager for manager in managers.values()
+                               if manager not in preexisting],
+                              key=lambda manager: manager.num)
+        new_figs[:] = [manager.canvas.figure for manager in new_managers]
+
+
 def _raise_on_image_difference(expected, actual, tol):
     __tracebackhide__ = True
 
@@ -178,10 +197,8 @@ def copy_baseline(self, baseline, extension):
                 f"{orig_expected_path}") from err
         return expected_fname
 
-    def compare(self, idx, baseline, extension, *, _lock=False):
+    def compare(self, fig, baseline, extension, *, _lock=False):
         __tracebackhide__ = True
-        fignum = plt.get_fignums()[idx]
-        fig = plt.figure(fignum)
 
         if self.remove_text:
             remove_ticks_and_titles(fig)
@@ -196,7 +213,12 @@ def compare(self, idx, baseline, extension, *, _lock=False):
         lock = (cbook._lock_path(actual_path)
                 if _lock else contextlib.nullcontext())
         with lock:
-            fig.savefig(actual_path, **kwargs)
+            try:
+                fig.savefig(actual_path, **kwargs)
+            finally:
+                # Matplotlib has an autouse fixture to close figures, but this
+                # makes things more convenient for third-party users.
+                plt.close(fig)
             expected_path = self.copy_baseline(baseline, extension)
             _raise_on_image_difference(expected_path, actual_path, self.tol)
 
@@ -235,7 +257,9 @@ def wrapper(*args, extension, request, **kwargs):
             img = _ImageComparisonBase(func, tol=tol, remove_text=remove_text,
                                        savefig_kwargs=savefig_kwargs)
             matplotlib.testing.set_font_settings_for_testing()
-            func(*args, **kwargs)
+
+            with _collect_new_figures() as figs:
+                func(*args, **kwargs)
 
             # If the test is parametrized in any way other than applied via
             # this decorator, then we need to use a lock to prevent two
@@ -252,11 +276,11 @@ def wrapper(*args, extension, request, **kwargs):
                 our_baseline_images = request.getfixturevalue(
                     'baseline_images')
 
-            assert len(plt.get_fignums()) == len(our_baseline_images), (
+            assert len(figs) == len(our_baseline_images), (
                 "Test generated {} images but there are {} baseline images"
-                .format(len(plt.get_fignums()), len(our_baseline_images)))
-            for idx, baseline in enumerate(our_baseline_images):
-                img.compare(idx, baseline, extension, _lock=needs_lock)
+                .format(len(figs), len(our_baseline_images)))
+            for fig, baseline in zip(figs, our_baseline_images):
+                img.compare(fig, baseline, extension, _lock=needs_lock)
 
         parameters = list(old_sig.parameters.values())
         if 'extension' not in old_sig.parameters:
@@ -427,11 +451,9 @@ def wrapper(*args, ext, request, **kwargs):
             try:
                 fig_test = plt.figure("test")
                 fig_ref = plt.figure("reference")
-                # Keep track of number of open figures, to make sure test
-                # doesn't create any new ones
-                n_figs = len(plt.get_fignums())
-                func(*args, fig_test=fig_test, fig_ref=fig_ref, **kwargs)
-                if len(plt.get_fignums()) > n_figs:
+                with _collect_new_figures() as figs:
+                    func(*args, fig_test=fig_test, fig_ref=fig_ref, **kwargs)
+                if figs:
                     raise RuntimeError('Number of open figures changed during '
                                        'test. Make sure you are plotting to '
                                        'fig_test or fig_ref, or if this is '
