Fits.py

Go to the documentation of this file.

from typing import Union, Optional
from io import StringIO
 
from robot.api import logger
import numpy as np
from astropy.io import fits
 
MaybeHdu = Union[fits.HDUList, str]
 
 
class Fits:
    """Robot library providing FITS related keywords"""
 
    ROBOT_LIBRARY_SCOPE = "GLOBAL"
 
    def _get_hdu(self, arg: MaybeHdu):
        if isinstance(arg, str):
            return fits.open(arg)
        return arg
 
    def fits_read_primary_hdu(
        self, path_or_hdu: MaybeHdu, with_checksum: Optional[bool] = None
    ):
        """Returns list of keywords from primary HDU as a list of pairs
        [["KEYWORD": "VALUE"], ...]
        """
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        if len(hdul) == 0:
            raise Exception("FITS file contains no HDUs")
        if with_checksum:
            return [list(tup) for tup in hdul[0].header.items()]
        redact = ["CHECKSUM", "DATASUM", "ARCFILE"]
        kws = [list(tup) for tup in hdul[0].header.items()]
        # Redact certain keywords to make testing easier
        for r in redact:
            for kw in kws:
                if r == kw[0]:
                    kw[1] = "XXXX"
        return kws
 
    def fits_get_hdu_ext_names(self, path_or_hdu: MaybeHdu):
        """Returns list of keywords from primary HDU as a list of pairs"""
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        if len(hdul) == 0:
            raise Exception("FITS file contains no HDUs")
        if len(hdul) > 1:
            # Note: keywords are tuples with (name, value, comment)
            return [(hdu.header.get("EXTNAME")) for hdu in hdul[1:]]
        else:
            return []
 
    def fits_create_empty(self):
        """"""
        hdul = fits.HDUList([fits.PrimaryHDU()])
        return hdul
 
    def fits_create_image(self):
        """"""
        hdul = fits.HDUList([fits.PrimaryHDU(np.ones((100, 100)))])
        return hdul
 
    def fits_get_info(self, path_or_hdu: MaybeHdu):
        s = StringIO()
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        hdul.info(output=s)
 
        return s.getvalue()
 
    def fits_set_keyword(
        self, path_or_hdu: MaybeHdu, *, name, value, comment="", hdu_num=0
    ):
        """"""
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        logger.info(
            "Setting keyword in hdu %s: %s=%s (%s)" % (hdu_num, name, value, comment)
        )
        hdul[hdu_num].header.set(name, value, comment)
        return hdul
 
    def fits_add_image_ext(self, path_or_hdu: MaybeHdu, *, name):
        """"""
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        logger.info("Setting extname %s" % name)
        hdul.append(fits.ImageHDU(np.ones((100, 100)), name=name))
        return hdul
 
    def fits_add_bintable_ext(self, path_or_hdu: MaybeHdu, *, name):
        """"""
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        c1 = fits.Column(name="a", array=np.array([1, 2]), format="K")
        c2 = fits.Column(name="b", array=np.array([4, 5]), format="K")
        c3 = fits.Column(name="c", array=np.array([7, 8]), format="K")
        ext = fits.BinTableHDU.from_columns([c1, c2, c3])
        logger.info("Setting extname %s" % name)
        ext.name = name
        hdul.append(ext)
        return hdul
 
    def fits_close(self, hdul: fits.HDUList):
        hdul.close()
 
    def fits_write_to(self, path_or_hdu: MaybeHdu, path: str):
        self._get_hdu_get_hdu_get_hdu(path_or_hdu).writeto(path, overwrite=True)
 
    def fits_verify(self, path_or_hdu: MaybeHdu):
        """
        Executes astropy verify() verify_checksum() and verify_datasum()
        """
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        hdul.verify("exception")
        mapping = {0: "FAILURE", 1: "SUCCESS", 2: "ABSENT"}
        for hdu in hdul:
            res = hdu.verify_checksum()
            if res != 1:
                raise AssertionError(
                    "FITS checksum verification failed. Expected '%s' but got '%s"
                    % (mapping[res], expected)
                )
            res = hdu.verify_datasum()
            if res != 1:
                raise AssertionError(
                    "FITS datsum verification failed. Expected '%s' but got '%s"
                    % (mapping[res], expected)
                )
 
    def fits_verify_checksum(
        self,
        path_or_hdu: MaybeHdu,
        hdu_num: Union[str, int],
        expected: Optional[str] = None,
    ):
        """Verify checksum"""
        hdu_num = int(hdu_num)
        mapping = {0: "FAILURE", 1: "SUCCESS", 2: "ABSENT"}
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        hdu = hdul[hdu_num]
        res = hdu.verify_checksum()
        res_str = mapping[res]
        if expected is not None:
            raise AssertionError(
                "FITS checksum verification failed. Expected '%s' but got '%s"
                % (res_str, expected)
            )
        return res_str
 
    def fits_verify_datasum(
        self,
        path_or_hdu: MaybeHdu,
        hdu_num: Union[str, int],
        expected: Optional[str] = None,
    ):
        """Verify datasum"""
        mapping = {0: "FAILURE", 1: "SUCCESS", 2: "ABSENT"}
        hdul = self._get_hdu_get_hdu_get_hdu(path_or_hdu)
        hdu = hdul[hdu_num]
        res = hdu.verify_datasum()
        res_str = mapping[res]
        if expected is not None:
            raise AssertionError(
                "FITS datasum verification failed. Expected '%s' but got '%s"
                % (res_str, expected)
            )
        return res_str