Merge pull request #477 from drdavella/extract-asdf-in-fits

Extract ASDF-in-FITS extension into pure ASDF file

CHANGES.rst

@@ -15,6 +15,9 @@
 - Remove hard dependency on Astropy. It is still required for testing, and for
   processing ASDF-in-FITS files. [#476]
 
+- Add command for extracting ASDF extension from ASDF-in-FITS file and
+  converting it to a pure ASDF file. [#477]
+
 1.4.0 (unreleased)
 ------------------
 

asdf/commands/__init__.py

@@ -1,10 +1,15 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 # -*- coding: utf-8 -*-
 
+import importlib
 
 from .exploded import *
 from .to_yaml import *
 from .defragment import *
 from .diff import *
 from .tags import *
 from .extension import *
+
+# Extracting ASDF-in-FITS files requires Astropy
+if importlib.util.find_spec('astropy'):
+    from .extract import *

asdf/commands/extract.py

@@ -0,0 +1,58 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+# -*- coding: utf-8 -*-
+
+"""
+Implementation of command for converting ASDF-in-FITS to standalone ASDF file.
+"""
+
+import sys
+
+import asdf
+from asdf import AsdfFile
+from asdf.fits_embed import AsdfInFits
+
+from .main import Command
+
+
+__all__ = ['extract_file']
+
+
+class FitsExtractor(Command): # pragma: no cover
+    """This class is the plugin implementation for the asdftool runner."""
+
+    @classmethod
+    def setup_arguments(cls, subparsers):
+        parser = subparsers.add_parser(str("extract"),
+            help="Extract ASDF extensions in ASDF-in-FITS files into pure ASDF files",
+            description="Extracts ASDF extensions into pure ASDF files.")
+
+        parser.add_argument(
+            'infile', action='store', type=str,
+            help="Name of ASDF-in-FITS file containing extension to be extracted")
+        parser.add_argument(
+            'outfile', action='store', type=str,
+            help="Name of new pure ASDF file containing extracted extension")
+
+        parser.set_defaults(func=cls.run)
+
+        return parser
+
+    @classmethod
+    def run(cls, args):
+        return extract_file(args.infile, args.outfile)
+
+
+def extract_file(input_file, output_file):
+    """Function for performing extraction from ASDF-in-FITS to pure ASDF."""
+
+    try:
+        with asdf.open(input_file) as ih:
+            if not isinstance(ih, AsdfInFits):
+                msg = "Given input file '{}' is not ASDF-in-FITS"
+                raise RuntimeError(msg.format(input_file))
+
+            with asdf.AsdfFile(ih.tree) as oh:
+                oh.write_to(output_file)
+
+    except (IOError, ValueError) as error:
+        raise RuntimeError(str(error))

asdf/commands/tests/test_extract.py

@@ -0,0 +1,47 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+# -*- coding: utf-8 -*-
+
+import os
+
+import numpy as np
+
+import pytest
+
+astropy = pytest.importorskip('astropy')
+from astropy.io.fits import HDUList, ImageHDU
+
+import asdf
+from asdf.fits_embed import AsdfInFits
+from asdf.tests.helpers import assert_tree_match
+
+from .. import extract
+
+
+def test_extract(tmpdir):
+    hdulist = HDUList()
+
+    image = ImageHDU(np.random.random((25, 25)))
+    hdulist.append(image)
+
+    tree = {
+        'some_words': 'These are some words',
+        'nested': {
+            'a': 100,
+            'b': 42
+        },
+        'list': [x for x in range(10)],
+        'image': image.data
+    }
+
+    asdf_in_fits = str(tmpdir.join('asdf.fits'))
+    with AsdfInFits(hdulist, tree) as aif:
+        aif.write_to(asdf_in_fits)
+
+    pure_asdf = str(tmpdir.join('extract.asdf'))
+    extract.extract_file(asdf_in_fits, pure_asdf)
+
+    assert os.path.exists(pure_asdf)
+
+    with asdf.open(pure_asdf) as af:
+        assert not isinstance(af, AsdfInFits)
+        assert_tree_match(tree, af.tree)

setup.cfg

@@ -13,7 +13,7 @@ minversion = 3.1
 norecursedirs = build docs/_build
 doctest_plus = enabled
 remote_data_strict = True
-open_files_ignore = test.fits
+open_files_ignore = test.fits asdf.fits
 # Account for both the astropy test runner case and the native pytest case
 asdf_schema_root = asdf-standard/schemas asdf/schemas
 asdf_schema_skip_names = asdf-schema-1.0.0 draft-01