Add back support for return HDUs from tests and comparing full FITS files

Author: astrofrog

pytest_arraydiff/plugin.py

@@ -32,21 +32,68 @@
 from functools import wraps
 
 import os
-import sys
+import abc
 import shutil
 import tempfile
 import warnings
 
+from six import add_metaclass
+from six.moves.urllib.request import urlopen
+
 import pytest
 import numpy as np
 
-if sys.version_info[0] == 2:
-    from urllib import urlopen
-else:
-    from urllib.request import urlopen
+
+@add_metaclass(abc.ABCMeta)
+class BaseDiff(object):
+
+    @abc.abstractstaticmethod
+    def read(filename):
+        """
+        Given a filename, return a data object.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractstaticmethod
+    def write(filename, data, **kwargs):
+        """
+        Given a filename and a data object (and optional keyword arguments),
+        write the data to a file.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractclassmethod
+    def compare(self, reference_file, test_file, atol=None, rtol=None):
+        """
+        Given a reference and test filename, compare the data to the specified
+        absolute (``atol``) and relative (``rtol``) tolerances.
+
+        Should return two arguments: a boolean indicating whether the data are
+        identical, and a string giving the full error message if not.
+        """
+        raise NotImplementedError()
+
+
+class SimpleArrayDiff(BaseDiff):
+
+    @classmethod
+    def compare(cls, reference_file, test_file, atol=None, rtol=None):
+
+        array_ref = cls.read(reference_file)
+        array_new = cls.read(test_file)
+
+        try:
+            np.testing.assert_allclose(array_ref, array_new, atol=atol, rtol=rtol)
+        except AssertionError as exc:
+            message = "\n\na: {0}".format(test_file) + '\n'
+            message += "b: {0}".format(reference_file) + '\n'
+            message += exc.args[0]
+            return False, message
+        else:
+            return True, ""
 
 
-class FITSDiff(object):
+class FITSDiff(BaseDiff):
 
     extension = 'fits'
 
@@ -56,12 +103,20 @@ def read(filename):
         return fits.getdata(filename)
 
     @staticmethod
-    def write(filename, array, **kwargs):
+    def write(filename, data, **kwargs):
         from astropy.io import fits
-        return fits.writeto(filename, array, **kwargs)
+        if isinstance(data, np.ndarray):
+            data = fits.PrimaryHDU(data)
+        return data.writeto(filename, **kwargs)
+
+    @classmethod
+    def compare(cls, reference_file, test_file, atol=None, rtol=None):
+        from astropy.io.fits.diff import FITSDiff
+        diff = FITSDiff(reference_file, test_file, tolerance=rtol)
+        return diff.identical, diff.report()
 
 
-class TextDiff(object):
+class TextDiff(SimpleArrayDiff):
 
     extension = 'txt'
 
@@ -70,10 +125,10 @@ def read(filename):
         return np.loadtxt(filename)
 
     @staticmethod
-    def write(filename, array, **kwargs):
+    def write(filename, data, **kwargs):
         if 'fmt' not in kwargs:
             kwargs['fmt'] = '%g'
-        return np.savetxt(filename, array, **kwargs)
+        return np.savetxt(filename, data, **kwargs)
 
 
 FORMATS = {}
@@ -219,17 +274,12 @@ def item_function_wrapper(*args, **kwargs):
                 baseline_file = os.path.abspath(os.path.join(result_dir, 'reference-' + filename))
                 shutil.copyfile(baseline_file_ref, baseline_file)
 
-                array_ref = FORMATS[file_format].read(baseline_file)
+                identical, msg = FORMATS[file_format].compare(baseline_file, test_image, atol=atol, rtol=rtol)
 
-                try:
-                    np.testing.assert_allclose(array_ref, array, atol=atol, rtol=rtol)
-                except AssertionError as exc:
-                    message = "\n\na: {0}".format(test_image) + '\n'
-                    message += "b: {0}".format(baseline_file) + '\n'
-                    message += exc.args[0]
-                    raise AssertionError(message)
-
-                shutil.rmtree(result_dir)
+                if identical:
+                    shutil.rmtree(result_dir)
+                else:
+                    raise Exception(msg)
 
             else:
 

setup.py

@@ -16,9 +16,10 @@
     name="pytest-arraydiff",
     description='pytest plugin to help with comparing array output from tests',
     long_description=long_description,
-    packages = ['pytest_arraydiff'],
+    packages=['pytest_arraydiff'],
+    install_requires=['numpy', 'six', 'pytest'],
     license='BSD',
     author='Thomas Robitaille',
     author_email='[email protected]',
-    entry_points = {'pytest11': ['pytest_arraydiff = pytest_arraydiff.plugin',]},
-)
+    entry_points={'pytest11': ['pytest_arraydiff = pytest_arraydiff.plugin']},
+    )

tests/test_pytest_arraydiff.py

@@ -23,6 +23,12 @@ def test_succeeds_func_fits():
     return np.arange(3 * 5).reshape((3, 5))
 
 
+@pytest.mark.array_compare(file_format='fits', reference_dir=reference_dir)
+def test_succeeds_func_fits_hdu():
+    from astropy.io import fits
+    return fits.PrimaryHDU(np.arange(3 * 5).reshape((3, 5)))
+
+
 class TestClass(object):
 
     @pytest.mark.array_compare(file_format='fits', reference_dir=reference_dir)