Add command to asdftool for removing ASDF extension from FITS file

CHANGES.rst

@@ -18,6 +18,8 @@
 - Add command for extracting ASDF extension from ASDF-in-FITS file and
   converting it to a pure ASDF file. [#477]
 
+- Add command for removing ASDF extension from ASDF-in-FITS file. [#480]
+
 1.4.0 (unreleased)
 ------------------
 

asdf/commands/__init__.py

@@ -13,3 +13,4 @@
 # Extracting ASDF-in-FITS files requires Astropy
 if importlib.util.find_spec('astropy'):
     from .extract import *
+    from .remove_hdu import *

asdf/commands/remove_hdu.py

@@ -0,0 +1,51 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+# -*- coding: utf-8 -*-
+
+"""
+Implementation of command for removing ASDF HDU from ASDF-in-FITS file.
+"""
+
+import sys
+
+from astropy.io import fits
+
+from .main import Command
+
+
+__all__ = ['remove_hdu']
+
+
+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("remove-hdu"),
+            help="Remove ASDF extension from ASDF-in-FITS file",
+            description="Removes ASDF extensions from ASDF-in-FITS files.")
+
+        parser.add_argument('infile', action='store', type=str,
+            help="Name of ASDF-in-FITS file containing extension to be removed")
+        parser.add_argument('outfile', action='store', type=str,
+            help="Name of new FITS output file")
+
+        parser.set_defaults(func=cls.run)
+
+        return parser
+
+    @classmethod
+    def run(cls, args):
+        return remove_hdu(args.infile, args.outfile)
+
+
+def remove_hdu(input_file, output_file):
+    """Function for removing ASDF HDU from ASDF-in-FITS files"""
+
+    try:
+        with fits.open(input_file) as hdulist:
+            hdulist.readall()
+            asdf_hdu = hdulist['ASDF']
+            hdulist.remove(asdf_hdu)
+            hdulist.writeto(output_file)
+    except (ValueError, KeyError) as error:
+        raise RuntimeError(str(error))

asdf/commands/tests/test_remove_hdu.py

@@ -0,0 +1,46 @@
+# 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 import fits
+
+from asdf.fits_embed import AsdfInFits
+from .. import remove_hdu
+
+
+def test_remove_hdu(tmpdir):
+    hdulist = fits.HDUList()
+
+    image = fits.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)
+
+    with fits.open(asdf_in_fits) as hdul:
+        assert 'ASDF' in hdul
+
+    new_fits = str(tmpdir.join('remove.fits'))
+    remove_hdu(asdf_in_fits, new_fits)
+
+    assert os.path.exists(new_fits)
+
+    with fits.open(new_fits) as hdul:
+        assert 'ASDF' not in hdul