diff --git a/csst_common/io.py b/csst_common/io.py index e97da4fc57f26e73ab8113982c3e65af2894f3d7..79c6296de03ba73eb6dad278d67a340092ddcefc 100644 --- a/csst_common/io.py +++ b/csst_common/io.py @@ -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 diff --git a/tests/test_fits_header_ops.py b/tests/test_fits_header_ops.py index 17e154e3a6b38b11815297c57a4b5bf91960d7eb..f9048b725a1aef27eacaf7a43c8a87ded8abe591 100644 --- a/tests/test_fits_header_ops.py +++ b/tests/test_fits_header_ops.py @@ -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"), + )