Added feature to copy to and extract from mpl_cache the baseline images

SidharthBansal · SidharthBansal · commit 15d0e9034e5b · 2020-09-07T17:12:06.000+05:30
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -17,7 +17,7 @@
 from matplotlib import ft2font
 from matplotlib import pyplot as plt
 from matplotlib import ticker
-
+from matplotlib.testing.compare import get_cache_dir
 from .compare import comparable_formats, compare_images, make_test_filename
 from .exceptions import ImageComparisonFailure
 
@@ -508,6 +508,7 @@ def _image_directories(func):
     doesn't exist.
     """
     module_path = Path(sys.modules[func.__module__].__file__)
+    cache_dir = Path(get_cache_dir()) / "baseline_images" / module_path.stem
     if func.__module__.startswith("matplotlib."):
         try:
             import matplotlib_baseline_images
@@ -528,7 +529,20 @@ def _image_directories(func):
             baseline_dir = (module_path.parent /
                             "baseline_images" /
                             module_path.stem)
-
+    extensions = ['**/*.svg', '**/*.eps', '**/*.png', '**/*.pdf']
+    for ext in extensions:
+        for baseline_img in Path(baseline_dir).glob(ext):
+            cache_dir_img = (cache_dir) / \
+                            (baseline_img).relative_to(baseline_dir)
+            if not cache_dir_img.exists():
+                cache_dir_img.parent.mkdir(parents=True, exist_ok=True)
+                try:
+                    shutil.copyfile(baseline_img, cache_dir_img)
+                except shutil.SameFileError:
+                    if not cache_dir_img.exists():
+                        raise
+                except FileNotFoundError:
+                    raise
     result_dir = Path().resolve() / "result_images" / module_path.stem
     result_dir.mkdir(parents=True, exist_ok=True)
-    return baseline_dir, result_dir
+    return cache_dir, result_dir
