Only close actual Matplotlib Figure objects

Author: astrofrog

CHANGES.md

@@ -1,10 +1,13 @@
 0.9 (unreleased)
 ----------------
 
+- Fix compatibility with Matplotlib 2.1. [#54]
+
 - Allow baseline_dir to be comma-separated URL list to allow mirrors to
   be specified. [#59]
 
-- Make sure figures get closed even if not running with the --mpl option.
+- Make sure figures get closed even if not running with the --mpl
+  option, and only close actual Matplotlib Figure objects. [#60]
 
 0.8 (2017-07-19)
 ----------------

pytest_mpl/plugin.py

@@ -155,6 +155,7 @@ def pytest_runtest_setup(self, item):
 
         import matplotlib
         import matplotlib.pyplot as plt
+        from matplotlib.figure import Figure
         from matplotlib.testing.compare import compare_images
         from matplotlib.testing.decorators import ImageComparisonTest as MplImageComparisonTest
         try:
@@ -223,7 +224,13 @@ def item_function_wrapper(*args, **kwargs):
                     test_image = os.path.abspath(os.path.join(result_dir, filename))
 
                     fig.savefig(test_image, **savefig_kwargs)
-                    plt.close(fig)
+
+                    # We only need to close actual Matplotlib figure objects. If
+                    # we are dealing with a figure-like object that provides
+                    # savefig but is not a real Matplotlib object, we shouldn't
+                    # try closing it here.
+                    if isinstance(fig, Figure):
+                        plt.close(fig)
 
                     # Find path to baseline image
                     if baseline_remote:
@@ -280,6 +287,7 @@ def pytest_runtest_setup(self, item):
             return
 
         import matplotlib.pyplot as plt
+        from matplotlib.figure import Figure
 
         original = item.function
 
@@ -291,7 +299,12 @@ def item_function_wrapper(*args, **kwargs):
             else:  # function
                 fig = original(*args, **kwargs)
 
-            plt.close(fig)
+            # We only need to close actual Matplotlib figure objects. If
+            # we are dealing with a figure-like object that provides
+            # savefig but is not a real Matplotlib object, we shouldn't
+            # try closing it here.
+            if isinstance(fig, Figure):
+                plt.close(fig)
 
         if item.cls is not None:
             setattr(item.cls, item.function.__name__, item_function_wrapper)