Move check_figures_equal out of decorators.py and back into testing.py

weiji14 · weiji14 · commit fb19edf7cc2b · 2020-09-04T09:33:52.000+12:00
diff --git a/pygmt/helpers/__init__.py b/pygmt/helpers/__init__.py
@@ -1,7 +1,7 @@
 """
 Functions, classes, decorators, and context managers to help wrap GMT modules.
 """
-from .decorators import check_figures_equal, fmt_docstring, kwargs_to_strings, use_alias
+from .decorators import fmt_docstring, use_alias, kwargs_to_strings
 from .tempfile import GMTTempFile, unique_name
 from .utils import (
     data_kind,
diff --git a/pygmt/helpers/decorators.py b/pygmt/helpers/decorators.py
@@ -5,16 +5,13 @@
 arguments, insert common text into docstrings, transform arguments to strings,
 etc.
 """
-import functools
-import inspect
-import os
 import textwrap
+import functools
 
 import numpy as np
-from matplotlib.testing.compare import compare_images
 
-from ..exceptions import GMTImageComparisonFailure, GMTInvalidInput
 from .utils import is_nonstr_iter
+from ..exceptions import GMTInvalidInput
 
 COMMON_OPTIONS = {
     "R": """\
@@ -406,98 +403,3 @@ def remove_bools(kwargs):
         else:
             new_kwargs[arg] = value
     return new_kwargs
-
-
-def check_figures_equal(*, tol=0.0, result_dir="result_images"):
-    """
-    Decorator for test cases that generate and compare two figures.
-
-    The decorated function must take two arguments, *fig_ref* and *fig_test*,
-    and draw the reference and test images on them. After the function
-    returns, the figures are saved and compared.
-
-    This decorator is practically identical to matplotlib's check_figures_equal
-    function, but adapted for PyGMT figures. See also the original code at
-    https://matplotlib.org/3.3.1/api/testing_api.html#
-    matplotlib.testing.decorators.check_figures_equal
-
-    Parameters
-    ----------
-    tol : float
-        The RMS threshold above which the test is considered failed.
-    result_dir : str
-        The directory where the figures will be stored.
-
-    Examples
-    --------
-
-    >>> import pytest
-    >>> @check_figures_equal()
-    ... def test_check_figures_equal(fig_ref, fig_test):
-    ...     fig_ref.basemap(projection="X5c", region=[0, 5, 0, 5], frame=True)
-    ...     fig_test.basemap(projection="X5c", region=[0, 5, 0, 5], frame="af")
-    >>> test_check_figures_equal()
-
-    >>> import shutil
-    >>> @check_figures_equal(result_dir="tmp_result_images")
-    ... def test_check_figures_unequal(fig_ref, fig_test):
-    ...     fig_ref.basemap(projection="X5c", region=[0, 5, 0, 5], frame=True)
-    ...     fig_test.basemap(projection="X5c", region=[0, 3, 0, 3], frame=True)
-    >>> with pytest.raises(GMTImageComparisonFailure):
-    ...     test_check_figures_unequal()
-    >>> shutil.rmtree(path="tmp_result_images")
-
-    """
-
-    def decorator(func):
-
-        os.makedirs(result_dir, exist_ok=True)
-        old_sig = inspect.signature(func)
-
-        def wrapper(*args, **kwargs):
-            try:
-                from ..figure import Figure  # pylint: disable=import-outside-toplevel
-
-                fig_ref = Figure()
-                fig_test = Figure()
-                func(*args, fig_ref=fig_ref, fig_test=fig_test, **kwargs)
-                ref_image_path = os.path.join(
-                    result_dir, func.__name__ + "-expected.png"
-                )
-                test_image_path = os.path.join(result_dir, func.__name__ + ".png")
-                fig_ref.savefig(ref_image_path)
-                fig_test.savefig(test_image_path)
-
-                # Code below is adapted for PyGMT, and is originally based on
-                # matplotlib.testing.decorators._raise_on_image_difference
-                err = compare_images(
-                    expected=ref_image_path,
-                    actual=test_image_path,
-                    tol=tol,
-                    in_decorator=True,
-                )
-                if err is None:  # Images are the same
-                    os.remove(ref_image_path)
-                    os.remove(test_image_path)
-                else:  # Images are not the same
-                    for key in ["actual", "expected", "diff"]:
-                        err[key] = os.path.relpath(err[key])
-                    raise GMTImageComparisonFailure(
-                        "images not close (RMS %(rms).3f):\n\t%(actual)s\n\t%(expected)s "
-                        % err
-                    )
-            finally:
-                del fig_ref
-                del fig_test
-
-        parameters = [
-            param
-            for param in old_sig.parameters.values()
-            if param.name not in {"fig_test", "fig_ref"}
-        ]
-        new_sig = old_sig.replace(parameters=parameters)
-        wrapper.__signature__ = new_sig
-
-        return wrapper
-
-    return decorator
diff --git a/pygmt/helpers/testing.py b/pygmt/helpers/testing.py
@@ -0,0 +1,105 @@
+"""
+Helper functions for testing.
+"""
+
+import inspect
+import os
+
+from matplotlib.testing.compare import compare_images
+
+from ..exceptions import GMTImageComparisonFailure
+from ..figure import Figure
+
+
+def check_figures_equal(*, tol=0.0, result_dir="result_images"):
+    """
+    Decorator for test cases that generate and compare two figures.
+
+    The decorated function must take two arguments, *fig_ref* and *fig_test*,
+    and draw the reference and test images on them. After the function
+    returns, the figures are saved and compared.
+
+    This decorator is practically identical to matplotlib's check_figures_equal
+    function, but adapted for PyGMT figures. See also the original code at
+    https://matplotlib.org/3.3.1/api/testing_api.html#
+    matplotlib.testing.decorators.check_figures_equal
+
+    Parameters
+    ----------
+    tol : float
+        The RMS threshold above which the test is considered failed.
+    result_dir : str
+        The directory where the figures will be stored.
+
+    Examples
+    --------
+
+    >>> import pytest
+    >>> import shutil
+
+    >>> @check_figures_equal(result_dir="tmp_result_images")
+    ... def test_check_figures_equal(fig_ref, fig_test):
+    ...     fig_ref.basemap(projection="X5c", region=[0, 5, 0, 5], frame=True)
+    ...     fig_test.basemap(projection="X5c", region=[0, 5, 0, 5], frame="af")
+    >>> test_check_figures_equal()
+
+    >>> @check_figures_equal(result_dir="tmp_result_images")
+    ... def test_check_figures_unequal(fig_ref, fig_test):
+    ...     fig_ref.basemap(projection="X5c", region=[0, 5, 0, 5], frame=True)
+    ...     fig_test.basemap(projection="X5c", region=[0, 3, 0, 3], frame=True)
+    >>> with pytest.raises(GMTImageComparisonFailure):
+    ...     test_check_figures_unequal()
+
+    >>> shutil.rmtree(path="tmp_result_images")  # cleanup folder if tests pass
+    """
+
+    def decorator(func):
+
+        os.makedirs(result_dir, exist_ok=True)
+        old_sig = inspect.signature(func)
+
+        def wrapper(*args, **kwargs):
+            try:
+                fig_ref = Figure()
+                fig_test = Figure()
+                func(*args, fig_ref=fig_ref, fig_test=fig_test, **kwargs)
+                ref_image_path = os.path.join(
+                    result_dir, func.__name__ + "-expected.png"
+                )
+                test_image_path = os.path.join(result_dir, func.__name__ + ".png")
+                fig_ref.savefig(ref_image_path)
+                fig_test.savefig(test_image_path)
+
+                # Code below is adapted for PyGMT, and is originally based on
+                # matplotlib.testing.decorators._raise_on_image_difference
+                err = compare_images(
+                    expected=ref_image_path,
+                    actual=test_image_path,
+                    tol=tol,
+                    in_decorator=True,
+                )
+                if err is None:  # Images are the same
+                    os.remove(ref_image_path)
+                    os.remove(test_image_path)
+                else:  # Images are not the same
+                    for key in ["actual", "expected", "diff"]:
+                        err[key] = os.path.relpath(err[key])
+                    raise GMTImageComparisonFailure(
+                        "images not close (RMS %(rms).3f):\n\t%(actual)s\n\t%(expected)s "
+                        % err
+                    )
+            finally:
+                del fig_ref
+                del fig_test
+
+        parameters = [
+            param
+            for param in old_sig.parameters.values()
+            if param.name not in {"fig_test", "fig_ref"}
+        ]
+        new_sig = old_sig.replace(parameters=parameters)
+        wrapper.__signature__ = new_sig
+
+        return wrapper
+
+    return decorator
diff --git a/pygmt/tests/test_grdimage.py b/pygmt/tests/test_grdimage.py
@@ -8,7 +8,7 @@
 from .. import Figure
 from ..datasets import load_earth_relief
 from ..exceptions import GMTInvalidInput
-from ..helpers import check_figures_equal
+from ..helpers.testing import check_figures_equal
 
 
 @pytest.fixture(scope="module", name="grid")
