Making both package and generation flag work simultaneously

SidharthBansal · SidharthBansal · commit cdfbf4e1c104 · 2020-08-23T19:33:23.000+05:30
diff --git a/lib/matplotlib/testing/conftest.py b/lib/matplotlib/testing/conftest.py
@@ -17,7 +17,7 @@ def pytest_configure(config):
         ("markers", "baseline_images: Compare output against references."),
         ("markers", "pytz: Tests that require pytz to be installed."),
         ("markers", "network: Tests that reach out to the network."),
-        ("markers", "generate_images: Flag to generate matplotlib baseline images."),
+        ("markers", "generate_images: Flag to generate baseline images."),
         ("filterwarnings", "error"),
     ]:
         config.addinivalue_line(key, value)
@@ -152,7 +152,9 @@ def pytest_addoption(parser):
 
 def pytest_collection_modifyitems(config, items):
     if config.getoption("--generate_images"):
-        skip_non_matplotlib_baseline_image_generation_tests = pytest.mark.skip(reason="No need to run non image generation tests")
+        skip_non_baseline_img_gen_tests \
+            = pytest.mark.skip(reason="No need to run non "
+                                      "image generation tests")
         for item in items:
-           if "generate_images" not in item.keywords:
-               item.add_marker(skip_non_matplotlib_baseline_image_generation_tests)
+            if "generate_images" not in item.keywords:
+                item.add_marker(skip_non_baseline_img_gen_tests)
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -6,15 +6,9 @@
 from pathlib import Path
 import shutil
 import string
-import subprocess
 import sys
 import unittest
 import warnings
-import pytest
-try:
-    from contextlib import nullcontext
-except ImportError:
-    from contextlib import ExitStack as nullcontext  # Py3.6.
 
 import matplotlib.style
 import matplotlib.units
@@ -205,7 +199,6 @@ def copy_baseline(self, baseline, extension):
                 f"{orig_expected_path}") from err
         return expected_fname
 
-
     def create_baseline(self, baseline, extension):
         src_path = self.result_dir / baseline
         orig_src_path = src_path.with_suffix(f'.{extension}')
@@ -224,12 +217,11 @@ def create_baseline(self, baseline, extension):
             raise ValueError("Failed to put the image in the right place")
         return dest_path
 
-
-    def compare(self, idx, baseline, extension, *, _lock=False, generating=False):
+    def compare(self, idx, baseline, extension, *, _lock=False,
+                generating=False):
         __tracebackhide__ = True
         fignum = plt.get_fignums()[idx]
         fig = plt.figure(fignum)
-
         if self.remove_text:
             remove_ticks_and_titles(fig)
 
@@ -248,8 +240,8 @@ def compare(self, idx, baseline, extension, *, _lock=False, generating=False):
                 self.create_baseline(baseline, extension)
             else:
                 expected_path = self.copy_baseline(baseline, extension)
-                _raise_on_image_difference(expected_path, actual_path, self.tol)
-
+                _raise_on_image_difference(expected_path, actual_path,
+                                           self.tol)
 
 
 def _pytest_image_comparison(baseline_images, extensions, tol,
@@ -307,7 +299,8 @@ def wrapper(*args, extension, request, **kwargs):
                 "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, generating=generate_images)
+                img.compare(idx, baseline, extension, _lock=needs_lock,
+                            generating=generate_images)
 
         parameters = list(old_sig.parameters.values())
         if 'extension' not in old_sig.parameters:
@@ -515,9 +508,34 @@ def _image_directories(func):
     doesn't exist.
     """
     module_path = Path(sys.modules[func.__module__].__file__)
-    baseline_dir = (module_path.parent /
-                      "baseline_images" /
-                       module_path.stem)
+    if func.__module__.startswith("matplotlib."):
+        try:
+            import matplotlib_baseline_images
+            baseline_dir = (Path(matplotlib_baseline_images.__file__).parent /
+                            module_path.stem)
+        except:
+            if (Path(__file__).parent / 'baseline_images').exists():
+                baseline_dir = (module_path.parent /
+                                "baseline_images" /
+                                module_path.stem)
+            else:
+                raise ImportError("Not able to import"
+                                  "matplotlib_baseline_images")
+    elif func.__module__.startswith("mpl_toolkits."):
+        try:
+            import mpl_toolkits_baseline_images
+            baseline_file = mpl_toolkits_baseline_images.__file__
+            baseline_dir = (Path(baseline_file).parent /
+                            module_path.stem)
+        except:
+            if (Path(mpl_toolkits_baseline_images.__file__).parent /
+                            module_path.stem).exists():
+                baseline_dir = (module_path.parent /
+                                "baseline_images" /
+                                module_path.stem)
+            else:
+                raise ImportError("Not able to import "
+                                  "mpl_toolkits_baseline_images")
     result_dir = Path().resolve() / "result_images" / module_path.stem
     result_dir.mkdir(parents=True, exist_ok=True)
     return baseline_dir, result_dir
diff --git a/lib/matplotlib/tests/__init__.py b/lib/matplotlib/tests/__init__.py
@@ -1 +1,14 @@
 # Check for the matplotlib_baseline_images.
+from pathlib import Path
+
+
+try:
+    import matplotlib_baseline_images
+except:
+    if not (Path(__file__).parent / 'baseline_images').exists():
+        raise ImportError(
+            'The baseline image directory does not exist. '
+            'This is most likely because the test data is not installed '
+            'and baseline image directory is absent. You may need to'
+            'install matplotlib_baseline_images to get the test data.'
+        )
diff --git a/lib/mpl_toolkits/tests/__init__.py b/lib/mpl_toolkits/tests/__init__.py
@@ -0,0 +1,12 @@
+from pathlib import Path
+
+try:
+    import matplotlib_baseline_images
+except:
+    if not (Path(__file__).parent / "baseline_images").exists():
+        raise ImportError(
+            'The baseline image directory does not exist. '
+            'This is most likely because the test data is not installed '
+            'and baseline image directory is absent. You may need to '
+            'install matplotlib_baseline_images to get the test data.'
+        )
diff --git a/lib/mpl_toolkits/tests/conftest.py b/lib/mpl_toolkits/tests/conftest.py
@@ -4,4 +4,3 @@
                                          pytest_unconfigure,
                                          pytest_collection_modifyitems,
                                          pytest_addoption)
-
diff --git a/tools/triage_tests.py b/tools/triage_tests.py
@@ -37,10 +37,22 @@
 # order to find the source, we need to search for a match in one of
 # these two places.
 
-BASELINE_IMAGES = [
-    Path('lib/matplotlib/tests/baseline_images'),
-    Path('lib/mpl_toolkits/tests/baseline_images'),
-]
+try:
+    import matplotlib_baseline_images
+    import mpl_toolkits_baseline_images
+    BASELINE_IMAGES = [
+        (Path(matplotlib_baseline_images.__file__).parent),
+        (Path(mpl_toolkits_baseline_images.__file__).parent),
+    ]
+except:
+    if not (Path('lib/matplotlib/tests/baseline_images')).exists()\
+    and not (Path('lib/mpl_toolkits/tests/baseline_images')).exists():
+        raise ImportError("Not able to import baseline images packages")
+    else:
+        BASELINE_IMAGES = [
+            Path('lib/matplotlib/tests/baseline_images'),
+            Path('lib/mpl_toolkits/tests/baseline_images'),
+        ]
 
 # Non-png image extensions
 
