Better handling of AsdfInFits

CHANGES.rst

@@ -8,6 +8,11 @@
 - Added a function ``is_asdf_file`` which inspects the input and
   returns ``True`` or ``False``. [#239]
 
+- The ``open`` method of ``AsdfInFits`` now accepts URIs and open file handles
+  in addition to HDULists. The ``open`` method of ``AsdfFile`` will now try to
+  parse the given URI or file handle as ``AsdfInFits`` if it is not obviously a
+  regular ASDF file. [#241]
+
 1.2.1(2016-11-07)
 -----------------
 

asdf/asdf.py

@@ -419,13 +419,11 @@ def _find_asdf_version_in_comments(cls, comments):
         return None
 
     @classmethod
-    def _open_impl(cls, self, fd, uri=None, mode='r',
+    def _open_asdf(cls, self, fd, uri=None, mode='r',
                    validate_checksums=False,
                    do_not_fill_defaults=False,
                    _get_yaml_content=False):
-        if not is_asdf_file(fd):
-            raise ValueError("Does not appear to be a ASDF file.")
-
+        """Attempt to populate AsdfFile data from file-like object"""
         fd = generic_io.get_file(fd, mode=mode, uri=uri)
         self._fd = fd
         header_line = fd.read_until(b'\r?\n', 2, "newline", include=True)
@@ -484,6 +482,30 @@ def _open_impl(cls, self, fd, uri=None, mode='r',
 
         return self
 
+    @classmethod
+    def _open_impl(cls, self, fd, uri=None, mode='r',
+                   validate_checksums=False,
+                   do_not_fill_defaults=False,
+                   _get_yaml_content=False):
+        """Attempt to open file-like object as either AsdfFile or AsdfInFits"""
+        if not is_asdf_file(fd):
+            try:
+                # TODO: this feels a bit circular, try to clean up. Also this
+                # introduces another dependency on astropy which may not be
+                # desireable.
+                from . import fits_embed
+                return fits_embed.AsdfInFits.open(fd, uri=uri,
+                            validate_checksums=validate_checksums,
+                            extensions=self._extensions)
+            except ValueError:
+                msg = "Input object does not appear to be ASDF file or " \
+                      "FITS ASDF extension"
+                raise ValueError(msg)
+        return cls._open_asdf(self, fd, uri=uri, mode=mode,
+                validate_checksums=validate_checksums,
+                do_not_fill_defaults=do_not_fill_defaults,
+                _get_yaml_content=_get_yaml_content)
+
     @classmethod
     def open(cls, fd, uri=None, mode='r',
              validate_checksums=False,
@@ -983,20 +1005,18 @@ def is_asdf_file(fd):
         return True
     elif isinstance(fd, generic_io.GenericFile):
         pass
-    elif isinstance(fd, io.IOBase):
+    else:
         try:
             fd = generic_io.get_file(fd, mode='r', uri=None)
+            if not isinstance(fd, io.IOBase):
+                to_close = True
         except ValueError:
             return False
-    else:
-        to_close = True
-        fd = generic_io.get_file(fd, mode='r', uri=None)
     asdf_magic = fd.read(5)
     if fd.seekable():
         fd.seek(0)
     if to_close:
         fd.close()
     if asdf_magic == constants.ASDF_MAGIC:
         return True
-    else:
-        return False
+    return False

asdf/fits_embed.py

@@ -16,6 +16,7 @@
 from . import asdf
 from . import block
 from . import util
+from . import generic_io
 
 try:
     from astropy.io import fits
@@ -146,27 +147,70 @@ def __init__(self, hdulist=None, tree=None, uri=None, extensions=None):
             tree=tree, uri=uri, extensions=extensions)
         self._blocks = _EmbeddedBlockManager(hdulist, self)
         self._hdulist = hdulist
+        self._close_hdulist = False
 
     def __exit__(self, type, value, traceback):
         super(AsdfInFits, self).__exit__(type, value, traceback)
+        if self._close_hdulist:
+            self._hdulist.close()
         self._tree = {}
 
     def close(self):
         super(AsdfInFits, self).close()
+        if self._close_hdulist:
+            self._hdulist.close()
         self._tree = {}
 
     @classmethod
-    def open(cls, hdulist, uri=None, validate_checksums=False, extensions=None):
+    def open(cls, fd, uri=None, validate_checksums=False, extensions=None):
+        """Creates a new AsdfInFits object based on given input data
+
+        Parameters
+        ----------
+        fd : FITS HDUList instance, URI string, or file-like object
+            May be an already opened instance of a FITS HDUList instance,
+            string ``file`` or ``http`` URI, or a Python file-like object.
+
+        uri : str, optional
+            The URI for this ASDF file.  Used to resolve relative
+            references against.  If not provided, will be
+            automatically determined from the associated file object,
+            if possible and if created from `AsdfFile.open`.
+
+        validate_checksums : bool, optional
+            If `True`, validate the blocks against their checksums.
+            Requires reading the entire file, so disabled by default.
+
+        extensions : list of AsdfExtension, optional
+            A list of extensions to the ASDF to support when reading
+            and writing ASDF files.  See `asdftypes.AsdfExtension` for
+            more information.
+        """
+        close_hdulist = False
+        if isinstance(fd, fits.hdu.hdulist.HDUList):
+            hdulist = fd
+        else:
+            file_obj = generic_io.get_file(fd)
+            try:
+                hdulist = fits.open(file_obj)
+                # Since we created this HDUList object, we need to be
+                # responsible for cleaning up upon close() or __exit__
+                close_hdulist = True
+            except IOError:
+                msg = "Failed to parse given file '{}'. Is it FITS?"
+                raise ValueError(msg.format(file_obj.uri))
+
         self = cls(hdulist, uri=uri, extensions=extensions)
+        self._close_hdulist = close_hdulist
 
         try:
             asdf_extension = hdulist[ASDF_EXTENSION_NAME]
         except (KeyError, IndexError, AttributeError):
+            # This means there is no ASDF extension
             return self
 
         buff = io.BytesIO(asdf_extension.data)
-
-        return cls._open_impl(self, buff, uri=uri, mode='r',
+        return cls._open_asdf(self, buff, uri=uri, mode='r',
                               validate_checksums=validate_checksums)
 
     def _update_asdf_extension(self, all_array_storage=None,

asdf/tests/test_fits_embed.py

@@ -21,10 +21,36 @@
 import pytest
 
 from .. import asdf
-
+from .. import open as asdf_open
 from .helpers import assert_tree_match
 
 
+def create_asdf_in_fits():
+    """Test fixture to create AsdfInFits object to use for testing"""
+    hdulist = fits.HDUList()
+    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
+    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
+    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
+
+    tree = {
+        'model': {
+            'sci': {
+                'data': hdulist[0].data,
+                'wcs': 'WCS info'
+            },
+            'dq': {
+                'data': hdulist[1].data,
+                'wcs': 'WCS info'
+            },
+            'err': {
+                'data': hdulist[2].data,
+                'wcs': 'WCS info'
+            }
+        }
+    }
+
+    return fits_embed.AsdfInFits(hdulist, tree)
+
 @pytest.mark.skipif('not HAS_ASTROPY')
 def test_embed_asdf_in_fits_file(tmpdir):
     hdulist = fits.HDUList()
@@ -68,47 +94,26 @@ def test_embed_asdf_in_fits_file(tmpdir):
 
 @pytest.mark.skipif('not HAS_ASTROPY')
 def test_embed_asdf_in_fits_file_anonymous_extensions(tmpdir):
-    hdulist = fits.HDUList()
-    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
-    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
-    hdulist.append(fits.ImageHDU(np.arange(512, dtype=np.float)))
+    # Write the AsdfInFits object out as a FITS file with ASDF extension
+    asdf_in_fits = create_asdf_in_fits()
+    asdf_in_fits.write_to(os.path.join(str(tmpdir), 'test.fits'))
 
-    tree = {
-        'model': {
-            'sci': {
-                'data': hdulist[0].data,
-                'wcs': 'WCS info'
-            },
-            'dq': {
-                'data': hdulist[1].data,
-                'wcs': 'WCS info'
-            },
-            'err': {
-                'data': hdulist[1].data,
-                'wcs': 'WCS info'
-            }
-        }
-    }
-
-    ff = fits_embed.AsdfInFits(hdulist, tree)
-    ff.write_to(os.path.join(str(tmpdir), 'test.fits'))
-
-    ff2 = asdf.AsdfFile(tree)
+    ff2 = asdf.AsdfFile(asdf_in_fits.tree)
     ff2.write_to(os.path.join(str(tmpdir), 'plain.asdf'))
 
-    with fits.open(os.path.join(str(tmpdir), 'test.fits')) as hdulist2:
-        assert len(hdulist2) == 4
-        assert [x.name for x in hdulist2] == ['PRIMARY', '', '', 'ASDF']
-        assert hdulist2['ASDF'].data.tostring().strip().endswith(b"...")
+    with fits.open(os.path.join(str(tmpdir), 'test.fits')) as hdulist:
+        assert len(hdulist) == 4
+        assert [x.name for x in hdulist] == ['PRIMARY', '', '', 'ASDF']
+        assert hdulist['ASDF'].data.tostring().strip().endswith(b"...")
 
-        with fits_embed.AsdfInFits.open(hdulist2) as ff2:
-            assert_tree_match(tree, ff2.tree)
+        with fits_embed.AsdfInFits.open(hdulist) as ff2:
+            assert_tree_match(asdf_in_fits.tree, ff2.tree)
 
             ff = asdf.AsdfFile(copy.deepcopy(ff2.tree))
             ff.write_to('test.asdf')
 
     with asdf.AsdfFile.open('test.asdf') as ff:
-        assert_tree_match(tree, ff.tree)
+        assert_tree_match(asdf_in_fits.tree, ff.tree)
 
 
 @pytest.mark.skipif('not HAS_ASTROPY')
@@ -135,17 +140,81 @@ def test_create_in_tree_first(tmpdir):
     hdulist.append(fits.ImageHDU(tree['model']['dq']['data']))
     hdulist.append(fits.ImageHDU(tree['model']['err']['data']))
 
-    ff = fits_embed.AsdfInFits(hdulist, tree)
-    ff.write_to(os.path.join(str(tmpdir), 'test.fits'))
+    with fits_embed.AsdfInFits(hdulist, tree) as ff:
+        ff.write_to(os.path.join(str(tmpdir), 'test.fits'))
 
-    ff2 = asdf.AsdfFile(tree)
-    ff2.write_to(os.path.join(str(tmpdir), 'plain.asdf'))
+    with asdf.AsdfFile(tree) as ff:
+        ff.write_to(os.path.join(str(tmpdir), 'plain.asdf'))
 
-    with asdf.AsdfFile.open(os.path.join(str(tmpdir), 'plain.asdf')) as ff3:
-        assert_array_equal(ff3.tree['model']['sci']['data'],
+    with asdf.AsdfFile.open(os.path.join(str(tmpdir), 'plain.asdf')) as ff:
+        assert_array_equal(ff.tree['model']['sci']['data'],
                            np.arange(512, dtype=np.float))
 
-    with fits.open(os.path.join(str(tmpdir), 'test.fits')) as hdulist:
-        with fits_embed.AsdfInFits.open(hdulist) as ff4:
-            assert_array_equal(ff4.tree['model']['sci']['data'],
-                               np.arange(512, dtype=np.float))
+    # This tests the changes that allow FITS files with ASDF extensions to be
+    # opened directly by the top-level AsdfFile.open API
+    with asdf_open(os.path.join(str(tmpdir), 'test.fits')) as ff:
+        assert_array_equal(ff.tree['model']['sci']['data'],
+                           np.arange(512, dtype=np.float))
+
+def compare_asdfs(asdf0, asdf1):
+    # Make sure the trees match
+    assert_tree_match(asdf0.tree, asdf1.tree)
+    # Compare the data blocks
+    for key in asdf0.tree['model'].keys():
+        assert_array_equal(
+            asdf0.tree['model'][key]['data'],
+            asdf1.tree['model'][key]['data'])
+
+@pytest.mark.skipif('not HAS_ASTROPY')
+def test_asdf_in_fits_open(tmpdir):
+    """Test the open method of AsdfInFits"""
+    tmpfile = os.path.join(str(tmpdir), 'test.fits')
+    # Write the AsdfInFits object out as a FITS file with ASDF extension
+    asdf_in_fits = create_asdf_in_fits()
+    asdf_in_fits.write_to(tmpfile)
+
+    # Test opening the file directly from the URI
+    with fits_embed.AsdfInFits.open(tmpfile) as ff:
+        compare_asdfs(asdf_in_fits, ff)
+
+    # Test open/close without context handler
+    ff = fits_embed.AsdfInFits.open(tmpfile)
+    compare_asdfs(asdf_in_fits, ff)
+    ff.close()
+
+    # Test reading in the file from an already-opened file handle
+    with open(tmpfile, 'rb') as handle:
+        with fits_embed.AsdfInFits.open(handle) as ff:
+            compare_asdfs(asdf_in_fits, ff)
+
+    # Test opening the file as a FITS file first and passing the HDUList
+    with fits.open(tmpfile) as hdulist:
+        with fits_embed.AsdfInFits.open(hdulist) as ff:
+            compare_asdfs(asdf_in_fits, ff)
+
+@pytest.mark.skipif('not HAS_ASTROPY')
+def test_asdf_open(tmpdir):
+    """Test the top-level open method of the asdf module"""
+    tmpfile = os.path.join(str(tmpdir), 'test.fits')
+    # Write the AsdfInFits object out as a FITS file with ASDF extension
+    asdf_in_fits = create_asdf_in_fits()
+    asdf_in_fits.write_to(tmpfile)
+
+    # Test opening the file directly from the URI
+    with asdf_open(tmpfile) as ff:
+        compare_asdfs(asdf_in_fits, ff)
+
+    # Test open/close without context handler
+    ff = asdf_open(tmpfile)
+    compare_asdfs(asdf_in_fits, ff)
+    ff.close()
+
+    # Test reading in the file from an already-opened file handle
+    with open(tmpfile, 'rb') as handle:
+        with asdf_open(handle) as ff:
+            compare_asdfs(asdf_in_fits, ff)
+
+    # Test opening the file as a FITS file first and passing the HDUList
+    with fits.open(tmpfile) as hdulist:
+        with asdf_open(hdulist) as ff:
+            compare_asdfs(asdf_in_fits, ff)