Merge pull request #17267 from QuLogic/testing-lock

timhoffm · web-flow · commit a975e48ebddd · 2020-05-06T01:19:59.000+02:00
Improve image comparison decorator
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -9,6 +9,10 @@
 import sys
 import unittest
 import warnings
+try:
+    from contextlib import nullcontext
+except ImportError:
+    from contextlib import ExitStack as nullcontext  # Py3.6.
 
 import matplotlib as mpl
 import matplotlib.style
@@ -200,7 +204,7 @@ def copy_baseline(self, baseline, extension):
                 f"{orig_expected_path}") from err
         return expected_fname
 
-    def compare(self, idx, baseline, extension):
+    def compare(self, idx, baseline, extension, *, _lock=False):
         __tracebackhide__ = True
         fignum = plt.get_fignums()[idx]
         fig = plt.figure(fignum)
@@ -214,10 +218,12 @@ def compare(self, idx, baseline, extension):
             kwargs.setdefault('metadata',
                               {'Creator': None, 'Producer': None,
                                'CreationDate': None})
-        fig.savefig(actual_path, **kwargs)
 
-        expected_path = self.copy_baseline(baseline, extension)
-        _raise_on_image_difference(expected_path, actual_path, self.tol)
+        lock = cbook._lock_path(actual_path) if _lock else nullcontext()
+        with lock:
+            fig.savefig(actual_path, **kwargs)
+            expected_path = self.copy_baseline(baseline, extension)
+            _raise_on_image_difference(expected_path, actual_path, self.tol)
 
 
 def _pytest_image_comparison(baseline_images, extensions, tol,
@@ -227,43 +233,66 @@ def _pytest_image_comparison(baseline_images, extensions, tol,
     Decorate function with image comparison for pytest.
 
     This function creates a decorator that wraps a figure-generating function
-    with image comparison code. Pytest can become confused if we change the
-    signature of the function, so we indirectly pass anything we need via the
-    `mpl_image_comparison_parameters` fixture and extra markers.
+    with image comparison code.
     """
     import pytest
 
     extensions = map(_mark_skip_if_format_is_uncomparable, extensions)
+    KEYWORD_ONLY = inspect.Parameter.KEYWORD_ONLY
 
     def decorator(func):
+        old_sig = inspect.signature(func)
+
         @functools.wraps(func)
-        # Parameter indirection; see docstring above and comment below.
-        @pytest.mark.usefixtures('mpl_image_comparison_parameters')
         @pytest.mark.parametrize('extension', extensions)
-        @pytest.mark.baseline_images(baseline_images)
-        # END Parameter indirection.
         @pytest.mark.style(style)
         @_checked_on_freetype_version(freetype_version)
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args, extension, request, **kwargs):
             __tracebackhide__ = True
+            if 'extension' in old_sig.parameters:
+                kwargs['extension'] = extension
+            if 'request' in old_sig.parameters:
+                kwargs['request'] = request
+
             img = _ImageComparisonBase(func, tol=tol, remove_text=remove_text,
                                        savefig_kwargs=savefig_kwargs)
             matplotlib.testing.set_font_settings_for_testing()
             func(*args, **kwargs)
 
-            # Parameter indirection:
-            # This is hacked on via the mpl_image_comparison_parameters fixture
-            # so that we don't need to modify the function's real signature for
-            # any parametrization. Modifying the signature is very very tricky
-            # and likely to confuse pytest.
-            baseline_images, extension = func.parameters
-
-            assert len(plt.get_fignums()) == len(baseline_images), (
+            # If the test is parametrized in any way other than applied via
+            # this decorator, then we need to use a lock to prevent two
+            # processes from touching the same output file.
+            needs_lock = any(
+                marker.args[0] != 'extension'
+                for marker in request.node.iter_markers('parametrize'))
+
+            if baseline_images is not None:
+                our_baseline_images = baseline_images
+            else:
+                # Allow baseline image list to be produced on the fly based on
+                # current parametrization.
+                our_baseline_images = request.getfixturevalue(
+                    'baseline_images')
+
+            assert len(plt.get_fignums()) == len(our_baseline_images), (
                 "Test generated {} images but there are {} baseline images"
-                .format(len(plt.get_fignums()), len(baseline_images)))
-            for idx, baseline in enumerate(baseline_images):
-                img.compare(idx, baseline, extension)
+                .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)
+
+        parameters = list(old_sig.parameters.values())
+        if 'extension' not in old_sig.parameters:
+            parameters += [inspect.Parameter('extension', KEYWORD_ONLY)]
+        if 'request' not in old_sig.parameters:
+            parameters += [inspect.Parameter("request", KEYWORD_ONLY)]
+        new_sig = old_sig.replace(parameters=parameters)
+        wrapper.__signature__ = new_sig
+
+        # Reach a bit into pytest internals to hoist the marks from our wrapped
+        # function.
+        new_marks = getattr(func, 'pytestmark', []) + wrapper.pytestmark
+        wrapper.pytestmark = new_marks
 
         return wrapper
 
@@ -398,13 +427,11 @@ def decorator(func):
                              f"function has the signature {old_sig}")
 
         @pytest.mark.parametrize("ext", extensions)
-        def wrapper(*args, **kwargs):
-            ext = kwargs['ext']
-            if 'ext' not in old_sig.parameters:
-                kwargs.pop('ext')
-            request = kwargs['request']
-            if 'request' not in old_sig.parameters:
-                kwargs.pop('request')
+        def wrapper(*args, ext, request, **kwargs):
+            if 'ext' in old_sig.parameters:
+                kwargs['ext'] = ext
+            if 'request' in old_sig.parameters:
+                kwargs['request'] = request
 
             file_name = "".join(c for c in request.node.name
                                 if c in ALLOWED_CHARS)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3150,15 +3150,10 @@ def test_hist_stacked_weighted():
     ax.hist((d1, d2), weights=(w1, w2), histtype="stepfilled", stacked=True)
 
 
-@image_comparison(['stem.png', 'stem.png'], style='mpl20', remove_text=True)
-def test_stem():
-    # Note, we don't use @pytest.mark.parametrize, because in parallel this
-    # might cause one process result to overwrite another's.
-    for use_line_collection in [True, False]:
-        _test_stem(use_line_collection)
-
-
-def _test_stem(use_line_collection):
+@pytest.mark.parametrize("use_line_collection", [True, False],
+                         ids=['w/ line collection', 'w/o line collection'])
+@image_comparison(['stem.png'], style='mpl20', remove_text=True)
+def test_stem(use_line_collection):
     x = np.linspace(0.1, 2 * np.pi, 100)
     args = (x, np.cos(x))
     # Label is a single space to force a legend to be drawn, but to avoid any
diff --git a/lib/mpl_toolkits/tests/test_axes_grid.py b/lib/mpl_toolkits/tests/test_axes_grid.py
@@ -12,16 +12,10 @@
 # The original version of this test relied on mpl_toolkits's slightly different
 # colorbar implementation; moving to matplotlib's own colorbar implementation
 # caused the small image comparison error.
-@image_comparison(['imagegrid_cbar_mode.png', 'imagegrid_cbar_mode.png'],
+@pytest.mark.parametrize("legacy_colorbar", [False, True])
+@image_comparison(['imagegrid_cbar_mode.png'],
                   remove_text=True, style='mpl20', tol=0.3)
-def test_imagegrid_cbar_mode_edge():
-    # Note, we don't use @pytest.mark.parametrize, because in parallel this
-    # might cause one process result to overwrite another's.
-    for legacy_colorbar in [False, True]:
-        _test_imagegrid_cbar_mode_edge(legacy_colorbar)
-
-
-def _test_imagegrid_cbar_mode_edge(legacy_colorbar):
+def test_imagegrid_cbar_mode_edge(legacy_colorbar):
     mpl.rcParams["mpl_toolkits.legacy_colorbar"] = legacy_colorbar
 
     X, Y = np.meshgrid(np.linspace(0, 6, 30), np.linspace(0, 6, 30))
