Commit 684956e5 authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

add reformat_header

parent 920aa478

csst_common/io.py

+18 −0
Original line number Diff line number Diff line
@@ -84,3 +84,21 @@ def append_header( 
                original_h.remove(card.keyword)

    original_h.extend(extended_h, bottom=True)

    return original_h





def reformat_header(h: fits.Header, strip=True, comment=None):

    h_copy = deepcopy(h)

    # strip

    if strip:

        h_copy.strip()

    # remove blanks, comments, and history

    h_copy.remove("", ignore_missing=True, remove_all=True)

    h_copy.remove("COMMENT", ignore_missing=True, remove_all=True)

    h_copy.remove("HISTORY", ignore_missing=True, remove_all=True)

    # add new comment

    assert len(h.cards) > 0

    first_key = list(h.keys())[0]

    h_copy.add_comment("=" * 72, before=first_key)

    h_copy.add_comment(comment, before=first_key)

    h_copy.add_comment("=" * 72, before=first_key)

    return h_copy

tests/test_fits_header_ops.py

+24 −1
Original line number Diff line number Diff line
@@ -11,7 +11,7 @@ import unittest 


from astropy.io import fits



from csst_common.io import append_header

from csst_common.io import append_header, reformat_header





class TestAppendHeader(unittest.TestCase):
@@ -61,3 +61,26 @@ class TestAppendHeader(unittest.TestCase): 
            ),

            "delete mode failed",

        )



    def test_reformat_header(self):

        h = fits.Header()

        h.add_comment("A")

        h.add_comment("B")

        h.add_comment("B")

        h.add_comment("C")

        h.add_comment("X")

        h.set("A", 1)

        h.set("SIMPLE", True)

        h.set("NAXIS1", 1)



        h_rfmt = reformat_header(h, strip=True, comment="WCS info")

        self.assertEqual(

            tuple(h_rfmt.keys()),

            ("COMMENT", "COMMENT", "COMMENT", "A", "NAXIS1"),

        )



        h_rfmt = reformat_header(h, strip=False, comment="WCS info")

        self.assertEqual(

            tuple(h_rfmt.keys()),

            ("COMMENT", "COMMENT", "COMMENT", "A", "SIMPLE", "NAXIS1"),

        )