HDU Header for components

glue/core/data_exporters/gridded_fits.py

@@ -9,6 +9,22 @@
 __all__ = []
 
 
+def make_component_header(component, header):
+    """
+    Function that extracts information from components
+    and adds it to the data header. The input header is
+    expected to come from Data.coords.header by default.
+    :param component: glue Component
+    :param header: astropy.io.fits.header.Header
+    :return:
+    """
+
+    # Add units information
+    header["BUNIT"] = component.units
+
+    return header
+
+
 @data_exporter(label='FITS (1 component/HDU)', extension=['fits', 'fit'])
 def fits_writer(filename, data, components=None):
     """
@@ -28,6 +44,8 @@ def fits_writer(filename, data, components=None):
     else:
         mask = None
 
+    data_header = data.coords.header if hasattr(data.coords, "header") else None
+
     from astropy.io import fits
 
     hdus = fits.HDUList()
@@ -50,7 +68,8 @@ def fits_writer(filename, data, components=None):
             values[~mask] = np.nan
 
         # TODO: special behavior for PRIMARY?
-        hdu = fits.ImageHDU(values, name=cid.label)
+        header = make_component_header(comp, data_header) if data_header else None
+        hdu = fits.ImageHDU(values, name=cid.label, header=header)
         hdus.append(hdu)
 
     try:

glue/core/data_factories/fits.py

@@ -114,14 +114,16 @@ def new_data(suffix=True):
             if is_image_hdu(hdu):
                 shape = hdu.data.shape
                 coords = coordinates_from_header(hdu.header)
+                units = hdu.header["BUNIT"] if "BUNIT" in hdu.header else None
                 if not auto_merge or has_wcs(coords):
                     data = new_data(suffix=len(hdulist) > 1)
                 else:
                     try:
                         data = groups[extension_by_shape[shape]]
                     except KeyError:
                         data = new_data(suffix=len(hdulist) > 1)
-                data.add_component(component=hdu.data,
+                component = Component.autotyped(hdu.data, units=units)
+                data.add_component(component=component,
                                    label=hdu_name)
             elif is_table_hdu(hdu):
                 # Loop through columns and make component list
@@ -205,7 +207,9 @@ def casalike_cube(filename, **kwargs):
         header = hdulist[0].header
     result.coords = coordinates_from_header(header)
     for i in range(array.shape[0]):
-        result.add_component(array[[i]], label='STOKES %i' % i)
+        units = header["BUNIT"] if "BUNIT" in header else None
+        component = Component.autotyped(array[[i]], units=units)
+        result.add_component(component, label='STOKES %i' % i)
     return result