From 02b10c6852c386493a450cca7b9a81712a5244cf Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Mon, 30 Oct 2023 11:05:39 +0000
Subject: [PATCH] Add tests to make sure that the healpix functions work with
 the full range of valid inputs/outputs (including e.g. APE-14 WCS) and update
 docstrings.

---
 reproject/healpix/high_level.py         | 25 +++++++++------
 reproject/healpix/tests/test_healpix.py | 42 +++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/reproject/healpix/high_level.py b/reproject/healpix/high_level.py
index 4df2c581f..f9ed1ed65 100644
--- a/reproject/healpix/high_level.py
+++ b/reproject/healpix/high_level.py
@@ -24,12 +24,14 @@ def reproject_from_healpix(
               second element is a `~astropy.coordinates.BaseCoordinateFrame`
               instance or a string alias for a coordinate frame.
 
-    output_projection : `~astropy.wcs.WCS` or `~astropy.io.fits.Header`
-        The output projection, which can be either a `~astropy.wcs.WCS`
-        or a `~astropy.io.fits.Header` instance.
+    output_projection : `~astropy.wcs.wcsapi.BaseLowLevelWCS` or `~astropy.wcs.wcsapi.BaseHighLevelWCS` or `~astropy.io.fits.Header`
+         The output projection, which can be either a
+         `~astropy.wcs.wcsapi.BaseLowLevelWCS`,
+         `~astropy.wcs.wcsapi.BaseHighLevelWCS`, or a `~astropy.io.fits.Header`
+         instance.
     shape_out : tuple, optional
-        If ``output_projection`` is a `~astropy.wcs.WCS` instance, the
-        shape of the output data should be specified separately.
+        If ``output_projection`` is a WCS instance, the shape of the output
+        data should be specified separately.
     hdu_in : int or str, optional
         If ``input_data`` is a FITS file, specifies the HDU to use.
         (the default HDU for HEALPIX data is 1, unlike with image files where
@@ -89,13 +91,18 @@ def reproject_to_healpix(
     input_data : object
         The input data to reproject. This can be:
 
-            * The name of a FITS file
+            * The name of a FITS file as a `str` or a `pathlib.Path` object
             * An `~astropy.io.fits.HDUList` object
-            * An image HDU object such as a `~astropy.io.fits.PrimaryHDU` or
-              `~astropy.io.fits.ImageHDU` instance
+            * An image HDU object such as a `~astropy.io.fits.PrimaryHDU`,
+              `~astropy.io.fits.ImageHDU`, or `~astropy.io.fits.CompImageHDU`
+              instance
             * A tuple where the first element is a `~numpy.ndarray` and the
-              second element is either a `~astropy.wcs.WCS` or a
+              second element is either a
+              `~astropy.wcs.wcsapi.BaseLowLevelWCS`,
+              `~astropy.wcs.wcsapi.BaseHighLevelWCS`, or a
               `~astropy.io.fits.Header` object
+            * An `~astropy.nddata.NDData` object from which the ``.data`` and
+              ``.wcs`` attributes will be used as the input data.
 
     coord_system_out : `~astropy.coordinates.BaseCoordinateFrame` or str
         The output coordinate system for the HEALPIX projection.
diff --git a/reproject/healpix/tests/test_healpix.py b/reproject/healpix/tests/test_healpix.py
index 8b68adfb6..84d4a4e7d 100644
--- a/reproject/healpix/tests/test_healpix.py
+++ b/reproject/healpix/tests/test_healpix.py
@@ -138,3 +138,45 @@ def test_reproject_invalid_order():
             os.path.join(DATA, "bayestar.fits.gz"), reference_header, order="bicubic"
         )
     assert exc.value.args[0] == "Only nearest-neighbor and bilinear interpolation are supported"
+
+
+def test_reproject_to_healpix_input_types(valid_celestial_input_data):
+    array_ref, wcs_in_ref, input_value, kwargs_in = valid_celestial_input_data
+
+    # Compute reference
+
+    healpix_data_ref, footprint_ref = reproject_to_healpix((array_ref, wcs_in_ref), "C", nside=64)
+
+    # Compute test
+
+    healpix_data_test, footprint_test = reproject_to_healpix(
+        input_value, "C", nside=64, **kwargs_in
+    )
+
+    # Make sure there are some valid values
+
+    assert np.sum(~np.isnan(healpix_data_ref)) == 3
+
+    np.testing.assert_allclose(healpix_data_ref, healpix_data_test)
+    np.testing.assert_allclose(footprint_ref, footprint_test)
+
+
+def test_reproject_from_healpix_output_types(valid_celestial_output_projections):
+    wcs_out_ref, shape_ref, output_value, kwargs_out = valid_celestial_output_projections
+
+    array_input = np.random.random(12 * 64**2)
+
+    # Compute reference
+
+    output_ref, footprint_ref = reproject_from_healpix(
+        (array_input, "C"), wcs_out_ref, nested=True, shape_out=shape_ref
+    )
+
+    # Compute test
+
+    output_test, footprint_test = reproject_from_healpix(
+        (array_input, "C"), output_value, nested=True, **kwargs_out
+    )
+
+    np.testing.assert_allclose(output_ref, output_test)
+    np.testing.assert_allclose(footprint_ref, footprint_test)