Merge pull request #5 from astrofrog/fits-hdu-support

Add back FITS HDU support

Author: astrofrog

README.rst

@@ -4,20 +4,21 @@ About
 -----
 
 This is a `py.test <http://pytest.org>`__ plugin to facilitate the
-generation and comparison of arrays produced during tests (this is a
+generation and comparison of data arrays produced during tests (this is a
 spin-off from
 `pytest-arraydiff <https://github.com/astrofrog/pytest-arraydiff>`__).
 
-The basic idea is that you can write a test that generates a Numpy
-array. You can then either run the tests in a mode to **generate**
-reference files from the arrays, or you can run the tests in
-**comparison** mode, which will compare the results of the tests to the
-reference ones within some tolerance.
+The basic idea is that you can write a test that generates a Numpy array (or
+other related objects depending on the format). You can then either run the
+tests in a mode to **generate** reference files from the arrays, or you can run
+the tests in **comparison** mode, which will compare the results of the tests to
+the reference ones within some tolerance.
 
 At the moment, the supported file formats for the reference files are:
 
--  The FITS format (requires `astropy <http://www.astropy.org>`__)
 -  A plain text-based format (baed on Numpy ``loadtxt`` output)
+-  The FITS format (requires `astropy <http://www.astropy.org>`__). With this
+   format, tests can return either a Numpy array for a FITS HDU object.
 
 For more information on how to write tests to do this, see the **Using**
 section below.
@@ -66,7 +67,7 @@ function returns a plain Numpy array:
     def test_succeeds():
         return np.arange(3 * 5 * 4).reshape((3, 5, 4))
 
-To generate the reference FITS files, run the tests with the
+To generate the reference data files, run the tests with the
 ``--arraydiff-generate-path`` option with the name of the directory
 where the generated files should be placed:
 
@@ -79,7 +80,7 @@ be interpreted as being relative to where you are running ``py.test``.
 Make sure you manually check the reference images to ensure they are
 correct.
 
-Once you are happy with the generated FITS files, you should move them
+Once you are happy with the generated data files, you should move them
 to a sub-directory called ``reference`` relative to the test files (this
 name is configurable, see below). You can also generate the baseline
 images directly in the right directory.

pytest_arraydiff/plugin.py

@@ -32,21 +32,78 @@
 from functools import wraps
 
 import os
-import sys
+import abc
 import shutil
 import tempfile
 import warnings
 
+from six import add_metaclass, PY2
+from six.moves.urllib.request import urlopen
+
 import pytest
 import numpy as np
 
-if sys.version_info[0] == 2:
-    from urllib import urlopen
+
+if PY2:
+    def abstractstaticmethod(func):
+        return func
+    def abstractclassmethod(func):
+        return func
 else:
-    from urllib.request import urlopen
+    abstractstaticmethod = abc.abstractstaticmethod
+    abstractclassmethod = abc.abstractclassmethod
+
+
+@add_metaclass(abc.ABCMeta)
+class BaseDiff(object):
+
+    @abstractstaticmethod
+    def read(filename):
+        """
+        Given a filename, return a data object.
+        """
+        raise NotImplementedError()
+
+    @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()
+
+    @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 +113,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 +135,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 +284,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='thomas.robitaille@gmail.com',
-    entry_points = {'pytest11': ['pytest_arraydiff = pytest_arraydiff.plugin',]},
-)
+    entry_points={'pytest11': ['pytest_arraydiff = pytest_arraydiff.plugin']},
+    )

tests/baseline/test_succeeds_func_fits_hdu.fits

@@ -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)