From ceabf4421b21c5260b8053dae29cfa6f8df97898 Mon Sep 17 00:00:00 2001
From: BO ZHANG <bozhang@nao.cas.cn>
Date: Mon, 25 Dec 2023 22:00:03 +0800
Subject: [PATCH] fixed bug in append_header

---
 csst_common/io.py             | 27 +++++++++++++++------------
 tests/test_fits_header_ops.py | 27 +++++++++++++++++++++------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/csst_common/io.py b/csst_common/io.py
index 5379d0a..b2cf06b 100644
--- a/csst_common/io.py
+++ b/csst_common/io.py
@@ -45,7 +45,7 @@ def append_header(
     """
     Append h2 to h1.
 
-    Append fits headers, taken into aacount duplicated keywords.
+    Append fits headers, taken into account duplicated keywords.
 
     Parameters
     ----------
@@ -70,20 +70,23 @@ def append_header(
     # copy data
     original_h = deepcopy(h1)
     extended_h = deepcopy(h2)
-    original_keys = original_h.keys()
-    ignored_keys = ["", "COMMENT", "HISTORY"]
+    original_keys = tuple(original_h.keys())
+    extended_keys = tuple(extended_h.keys())
+    ignored_keys = ("", "COMMENT", "HISTORY")
     assert duplicates in ("delete", "update")
     if duplicates == "update":
-        for card in extended_h.cards:
-            if card.keyword not in ignored_keys and card.keyword in original_keys:
-                print(f"Update existing key *{card.keyword} in original fits.Header*")
-                original_h.set(card.keyword, card.value, card.comment)
-                extended_h.remove(card.keyword)
+        for k in extended_keys:
+            # print(card.keyword)
+            if k not in ignored_keys and k in original_keys:
+                print(f"Update existing key *{k} in original fits.Header*")
+                original_h.set(k, extended_h[k], extended_h.comments[k])
+                extended_h.remove(k)
     elif duplicates == "delete":
-        for card in extended_h.cards:
-            if card.keyword not in ignored_keys and card.keyword in original_keys:
-                print(f"Delete existing key *{card.keyword} in original fits.Header*")
-                original_h.remove(card.keyword)
+        for k in extended_keys:
+            # print(card.keyword)
+            if k not in ignored_keys and k in original_keys:
+                print(f"Delete existing key *{k} in original fits.Header*")
+                original_h.remove(k)
     original_h.extend(extended_h, bottom=True)
     return original_h
 
diff --git a/tests/test_fits_header_ops.py b/tests/test_fits_header_ops.py
index b0488a4..3f0d641 100644
--- a/tests/test_fits_header_ops.py
+++ b/tests/test_fits_header_ops.py
@@ -21,34 +21,44 @@ class TestFitsHeaderOps(unittest.TestCase):
         h2 = fits.Header()
 
         h1.set("A", 1, "comment")
-        h1.set("B", 2, "comment")
+        h1.set("B", 1, "comment")
+        h1.set("C", 1, "comment")
         h1.add_comment("=" * 72, before="A")
         h1.add_comment("one", before="A")
         h1.add_comment("=" * 72, before="A")
 
-        h2.set("B", 3, "comment")
-        h2.set("C", 4, "comment")
+        h2.set("B", 2, "comment")
+        h2.set("C", 2, "comment")
+        h2.set("D", 2, "comment")
         h2.add_comment("=" * 72, before="B")
         h2.add_comment("another", before="B")
         h2.add_comment("=" * 72, before="B")
 
+        h_update = append_header(h1, h2, duplicates="update")
         self.assertEqual(
-            tuple(append_header(h1, h2, duplicates="update").keys()),
+            tuple(h_update.keys()),
             (
                 "COMMENT",
                 "COMMENT",
                 "COMMENT",
                 "A",
                 "B",
+                "C",
                 "COMMENT",
                 "COMMENT",
                 "COMMENT",
-                "C",
+                "D",
             ),
             "update mode failed",
         )
+        self.assertEqual(h_update["A"], 1)
+        self.assertEqual(h_update["B"], 2)
+        self.assertEqual(h_update["C"], 2)
+        self.assertEqual(h_update["D"], 2)
+
+        h_delete = append_header(h1, h2, duplicates="delete")
         self.assertEqual(
-            tuple(append_header(h1, h2, duplicates="delete").keys()),
+            tuple(h_delete.keys()),
             (
                 "COMMENT",
                 "COMMENT",
@@ -59,9 +69,14 @@ class TestFitsHeaderOps(unittest.TestCase):
                 "COMMENT",
                 "B",
                 "C",
+                "D",
             ),
             "delete mode failed",
         )
+        self.assertEqual(h_delete["A"], 1)
+        self.assertEqual(h_delete["B"], 2)
+        self.assertEqual(h_delete["C"], 2)
+        self.assertEqual(h_delete["D"], 2)
 
     def test_reformat_header(self):
         h = fits.Header()
-- 
GitLab