add reformat_header

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

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"),
+        )