testing: implement image_comparison decorator

astraw · astraw · commit e2f8829baa13 · 2009-09-06T01:44:09.000Z
svn path=/trunk/matplotlib/; revision=7652
diff --git a/lib/matplotlib/testing/compare.py b/lib/matplotlib/testing/compare.py
@@ -72,7 +72,7 @@ def compare_float( expected, actual, relTol = None, absTol = None ):
       return None
 
 #-----------------------------------------------------------------------
-def compare_images( expected, actual, tol ):
+def compare_images( expected, actual, tol, in_decorator=False ):
    '''Compare two image files - not the greatest, but fast and good enough.
 
    = EXAMPLE
@@ -87,6 +87,8 @@ def compare_images( expected, actual, tol ):
    - actual    The filename of the actual image.
    - tol       The tolerance (a unitless float).  This is used to
                determinte the 'fuzziness' to use when comparing images.
+   - in_decorator If called from image_comparison decorator, this should be
+               True. (default=False)
    '''
 
    try:
@@ -113,11 +115,21 @@ def compare_images( expected, actual, tol ):
 
    if ( (rms / 10000.0) <= tol ):
       return None
-   else:
-      diff_image = os.path.join(os.path.dirname(actual),
-                                'failed-diff-'+os.path.basename(actual))
-      save_diff_image( expected, actual, diff_image )
 
+   diff_image = os.path.join(os.path.dirname(actual),
+                             'failed-diff-'+os.path.basename(actual))
+   save_diff_image( expected, actual, diff_image )
+
+   if in_decorator:
+      results = dict(
+         rms = rms,
+         expected = str(expected),
+         actual = str(actual),
+         diff = str(diff_image),
+         )
+      return results
+   else:
+      # old-style call from mplTest directory
       msg = "  Error: Image files did not match.\n"       \
             "  RMS Value: " + str( rms / 10000.0 ) + "\n" \
             "  Expected:\n    " + str( expected ) + "\n"  \
@@ -130,6 +142,8 @@ def save_diff_image( expected, actual, output ):
    from PIL import Image
    expectedImage = np.array(Image.open( expected ).convert("RGB")).astype(np.float)
    actualImage = np.array(Image.open( actual ).convert("RGB")).astype(np.float)
+   assert expectedImage.ndim==expectedImage.ndim
+   assert expectedImage.shape==expectedImage.shape
    absDiffImage = abs(expectedImage-actualImage)
    # expand differences in luminance domain
    absDiffImage *= 10
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -1,6 +1,9 @@
 from matplotlib.testing.noseclasses import KnownFailureTest, \
-     KnownFailureDidNotFailTest
+     KnownFailureDidNotFailTest, ImageComparisonFailure
 import sys
+import nose
+from matplotlib.cbook import get_sample_data
+from matplotlib.testing.compare import compare_images
 
 def knownfailureif(fail_condition, msg=None):
     # based on numpy.testing.dec.knownfailureif
@@ -24,3 +27,23 @@ def failer(*args, **kwargs):
             return result
         return nose.tools.make_decorator(f)(failer)
     return known_fail_decorator
+
+def image_comparison(baseline_images=None, tol=1e-3):
+    if baseline_images is None:
+        raise ValueError('baseline_images must be specified')
+    def compare_images_decorator(func):
+        def decorated_compare_images(*args,**kwargs):
+            result = func(*args,**kwargs)
+            for fname in baseline_images:
+                actual = fname
+                expected = get_sample_data('test_baseline_%s'%fname,
+                                           asfileobj=False)
+                err = compare_images( expected, actual, tol,
+                                      in_decorator=True )
+                if err:
+                    raise ImageComparisonFailure(
+                        'images not close: %(actual)s vs. %(expected)s '
+                        '(RMS %(rms).3f)'%err)
+            return result
+        return nose.tools.make_decorator(func)(decorated_compare_images)
+    return compare_images_decorator
diff --git a/lib/matplotlib/testing/noseclasses.py b/lib/matplotlib/testing/noseclasses.py
@@ -9,6 +9,9 @@ class KnownFailureDidNotFailTest(Exception):
     '''Raise this exception to mark a test should have failed but did not.'''
     pass
 
+class ImageComparisonFailure(Exception):
+    '''Raise this exception to mark a test as a comparison between two images.'''
+
 class KnownFailure(ErrorClassPlugin):
     '''Plugin that installs a KNOWNFAIL error class for the
     KnownFailureClass exception.  When KnownFailureTest is raised,
