Make sure running the tests without --mpl still results in figures getting closed to avoid memory issues.

Author: astrofrog

.travis.yml

@@ -21,7 +21,7 @@ env:
     - PYTHON_VERSION=3.6 MATPLOTLIB_VERSION=2.0
 
     # The following build is meant to check that dependencies get set up correctly, but currently
-    # the image tests fail, which should be investigated. 
+    # the image tests fail, which should be investigated.
     # - PYTHON_VERSION=3.5 CONDA_DEPENDENCIES="coverage freetype libpng"
 
 install:
@@ -37,6 +37,9 @@ install:
 script:
    - python -c 'import pytest_mpl.plugin'
    - pytest -vv --mpl --cov pytest_mpl tests
+   # Make sure that the tests run ok even without the --mpl option (we close
+   # figures anyway in this case)
+   - pytest -vv --cov pytest_mpl --cov-append tests
    - python setup.py check --restructuredtext
 
 after_success:

CHANGES.md

@@ -4,6 +4,8 @@
 - 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.
+
 0.8 (2017-07-19)
 ----------------
 

pytest_mpl/plugin.py

@@ -34,6 +34,7 @@
 import os
 import sys
 import shutil
+import inspect
 import tempfile
 import warnings
 from distutils.version import LooseVersion
@@ -42,7 +43,7 @@
 
 if sys.version_info[0] == 2:
     from urllib import urlopen
-    string_types = basestring
+    string_types = basestring  # noqa
 else:
     from urllib.request import urlopen
     string_types = str
@@ -70,7 +71,7 @@ def _download_file(baseline, filename):
 
 
 def pytest_addoption(parser):
-    group = parser.getgroup("general")
+    group = parser.getgroup("matplotlib image comparison")
     group.addoption('--mpl', action='store_true',
                     help="Enable comparison of matplotlib figures to reference files")
     group.addoption('--mpl-generate-path',
@@ -117,6 +118,10 @@ def pytest_configure(config):
                                                       generate_dir=generate_dir,
                                                       results_dir=results_dir))
 
+    else:
+
+        config.pluginmanager.register(FigureCloser(config))
+
 
 @contextlib.contextmanager
 def switch_backend(backend):
@@ -143,6 +148,11 @@ def __init__(self, config, baseline_dir=None, generate_dir=None, results_dir=Non
 
     def pytest_runtest_setup(self, item):
 
+        compare = item.keywords.get('mpl_image_compare')
+
+        if compare is None:
+            return
+
         import matplotlib
         import matplotlib.pyplot as plt
         from matplotlib.testing.compare import compare_images
@@ -154,11 +164,6 @@ def pytest_runtest_setup(self, item):
 
         MPL_LT_15 = LooseVersion(matplotlib.__version__) < LooseVersion('1.5')
 
-        compare = item.keywords.get('mpl_image_compare')
-
-        if compare is None:
-            return
-
         tolerance = compare.kwargs.get('tolerance', 2)
         savefig_kwargs = compare.kwargs.get('savefig_kwargs', {})
         style = compare.kwargs.get('style', 'classic')
@@ -188,7 +193,6 @@ def item_function_wrapper(*args, **kwargs):
             with plt.style.context(style), switch_backend(backend):
 
                 # Run test and get figure object
-                import inspect
                 if inspect.ismethod(original):  # method
                     # In some cases, for example if setup_method is used,
                     # original appears to belong to an instance of the test
@@ -257,3 +261,39 @@ def item_function_wrapper(*args, **kwargs):
             setattr(item.cls, item.function.__name__, item_function_wrapper)
         else:
             item.obj = item_function_wrapper
+
+
+class FigureCloser(object):
+    """
+    This is used in place of ImageComparison when the --mpl option is not used,
+    to make sure that we still close figures returned by tests.
+    """
+
+    def __init__(self, config):
+        self.config = config
+
+    def pytest_runtest_setup(self, item):
+
+        compare = item.keywords.get('mpl_image_compare')
+
+        if compare is None:
+            return
+
+        import matplotlib.pyplot as plt
+
+        original = item.function
+
+        @wraps(item.function)
+        def item_function_wrapper(*args, **kwargs):
+
+            if inspect.ismethod(original):  # method
+                fig = original.__func__(*args, **kwargs)
+            else:  # function
+                fig = original(*args, **kwargs)
+
+            plt.close(fig)
+
+        if item.cls is not None:
+            setattr(item.cls, item.function.__name__, item_function_wrapper)
+        else:
+            item.obj = item_function_wrapper