Repository: sunpy/sunpy

Issue:
Allow passing keyword arguments to the HDU object used in GenericMap.save
Currently we allow specifying the HDU object to `GenericMap.save` (which ends up in `sunpy.io.fits.write`), but not any way of specifying arguments to this HDU, for example the `quantize_level=` kwarg to the `fits.CompImageHDU`.

The offending line is here: https://github.com/sunpy/sunpy/blob/0757ca6f8b/sunpy/io/fits.py#L185

The difficulty here is we are already using `**kwargs` as arguments to `fits.writeto`, so we either need a `hdu_kwargs=` argument (which doesn't work well with other writers) or we need to work out a way of allowing passing *instances* of a HDU object into `write`. Which would mean something like: `aia.save("test.fits", hdu_type=fits.CompImageHDU(quantize_level=32))` would work.

The only worry with that idea to me is what happens if the HDU instance already has data or a header set?


Discussion:

Resolving patch:
diff --git a/sunpy/io/fits.py b/sunpy/io/fits.py
--- a/sunpy/io/fits.py
+++ b/sunpy/io/fits.py
@@ -161,9 +161,11 @@ def write(fname, data, header, hdu_type=None, **kwargs):
         n-dimensional data array.
     header : `dict`
         A header dictionary.
-    hdu_type: `None` or `~astropy.io.fits.CompImageHDU`
-        `None` will return a normal FITS files.
-        `~astropy.io.fits.CompImageHDU` will rice compress the FITS file.
+    hdu_type : `~astropy.io.fits.hdu.base.ExtensionHDU` instance or class, optional
+        By default, a FITS file is written with the map in its primary HDU.
+        If a type is given, a new HDU of this type will be created.
+        If a HDU instance is given, its data and header will be updated from the map.
+        Then that HDU instance will be written to the file.
     kwargs :
         Additional keyword arguments are given to
         `~astropy.io.fits.HDUList.writeto`.
@@ -183,7 +185,18 @@ def write(fname, data, header, hdu_type=None, **kwargs):
     if not hdu_type:
         hdu_type = fits.PrimaryHDU
 
-    hdu = hdu_type(data=data, header=fits_header)
+    if isinstance(hdu_type, (fits.PrimaryHDU, fits.hdu.base.ExtensionHDU)):
+        hdu = hdu_type  # HDU already initialised
+
+        # Merge `header` into HDU's header
+        # Values in `header` take priority, including cards such as
+        # 'SIMPLE' and 'BITPIX'.
+        hdu.header.extend(fits_header, strip=False, update=True)
+
+        # Set the HDU's data
+        hdu.data = data
+    else:
+        hdu = hdu_type(data=data, header=fits_header)
 
     if not isinstance(hdu, fits.PrimaryHDU):
         hdul = fits.HDUList([fits.PrimaryHDU(), hdu])
diff --git a/sunpy/map/mapbase.py b/sunpy/map/mapbase.py
--- a/sunpy/map/mapbase.py
+++ b/sunpy/map/mapbase.py
@@ -1378,9 +1378,11 @@ def save(self, filepath, filetype='auto', **kwargs):
             Location to save file to.
         filetype : str
             'auto' or any supported file extension.
-        hdu_type: None, `~astropy.io.fits.CompImageHDU`
-            `None` will return a normal FITS file.
-            `~astropy.io.fits.CompImageHDU` will rice compress the FITS file.
+        hdu_type : `~astropy.io.fits.hdu.base.ExtensionHDU` instance or class, optional
+            By default, a FITS file is written with the map in its primary HDU.
+            If a type is given, a new HDU of this type will be created.
+            If a HDU instance is given, its data and header will be updated from the map.
+            Then that HDU instance will be written to the file.
         kwargs :
             Any additional keyword arguments are passed to
             `~sunpy.io.write_file`.


Test patch:
diff --git a/sunpy/io/tests/test_fits.py b/sunpy/io/tests/test_fits.py
--- a/sunpy/io/tests/test_fits.py
+++ b/sunpy/io/tests/test_fits.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from collections import OrderedDict
 
+import numpy as np
 import pytest
 
 import astropy.io.fits as fits
@@ -83,6 +84,48 @@ def test_simple_write_compressed(tmpdir):
         assert isinstance(hdul[1], fits.CompImageHDU)
 
 
+def test_simple_write_compressed_difftypeinst(tmpdir):
+    # `hdu_type=fits.CompImageHDU` and `hdu_type=fits.CompImageHDU()`
+    # should produce identical FITS files
+    data, header = sunpy.io.fits.read(AIA_171_IMAGE)[0]
+    outfile_type = str(tmpdir / "test_type.fits")
+    outfile_inst = str(tmpdir / "test_inst.fits")
+    sunpy.io.fits.write(outfile_type, data, header, hdu_type=fits.CompImageHDU)
+    sunpy.io.fits.write(outfile_inst, data, header, hdu_type=fits.CompImageHDU())
+    assert fits.FITSDiff(outfile_type, outfile_inst, ignore_comments=['PCOUNT']).identical
+
+
+@pytest.mark.parametrize(
+    'kwargs, should_fail',
+    [({}, False),
+     ({'quantize_level': -32}, True)]
+)
+def test_simple_write_compressed_instance(tmpdir, kwargs, should_fail):
+    data, header = sunpy.io.fits.read(AIA_171_IMAGE)[0]
+    outfile = tmpdir / "test.fits"
+
+    # Ensure HDU instance is used correctly
+    hdu = fits.CompImageHDU(data=np.array([0.]), **kwargs)
+    hdu.header['HELLO'] = 'world'  # should be in the written file
+    hdu.header['TELESCOP'] = 'other'  # should be replaced with 'SDO/AIA'
+    hdu.header['NAXIS'] = 5  # should be replaced with 2
+    sunpy.io.fits.write(str(outfile), data, header, hdu_type=hdu)
+    assert outfile.exists()
+    with fits.open(str(outfile)) as hdul:
+        assert len(hdul) == 2
+        assert isinstance(hdul[1], fits.CompImageHDU)
+        assert hdul[1].header['HELLO'] == 'world'
+        assert hdul[1].header['TELESCOP'] == 'SDO/AIA'
+        assert hdul[1].header['NAXIS'] == 2
+        data_preserved = hdul[1].data == pytest.approx(data, abs=10)
+        print(np.abs(hdul[1].data - data).max())
+        print(kwargs)
+        if should_fail:  # high compression setting preserved
+            assert not data_preserved
+        else:
+            assert data_preserved
+
+
 def test_write_with_metadict_header_astropy(tmpdir):
     with fits.open(AIA_171_IMAGE) as fits_file:
         data, header = fits_file[0].data, fits_file[0].header