Repository: astropy/astropy

Issue:
Fits Header does not preserve subclasses in copying and slicing.
I have a subclass of `io.fits.Header` and noticed that the copying and slicing does not by default return instances of my subclasses but rather of `Header`. This is because the class is hardcoded in those methods. In principle, this is incorrect, and correcting it is relatively straightforward [1], except that this breaks the `CompImageHeader` subclass since that cannot deal properly with its own class initializer being used. Which puts it slightly out of reach of a quick fix, and makes it more part of a cleanup of `CompImageHeader` that the comments for that class already refer to as being needed...

[1] https://github.com/astropy/astropy/compare/master...mhvk:io-fits-header-subclass-preservation?expand=1

Discussion:

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 0bd78b86ac..e3020fffc2 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -205,6 +205,10 @@ astropy.io.fits
 
 - Masked column handling has changed, see ``astropy.table`` entry below. [#8789]
 
+- ``io.fits.Header`` has been made safe for subclasses for copying and slicing.
+  As a result of this change, the private subclass ``CompImageHeader`` now always
+  should be passed an explicit ``image_header``. [#9229]
+
 astropy.io.registry
 ^^^^^^^^^^^^^^^^^^^
 
diff --git a/astropy/io/fits/hdu/compressed.py b/astropy/io/fits/hdu/compressed.py
index f288a8ae52..d1a922e72e 100644
--- a/astropy/io/fits/hdu/compressed.py
+++ b/astropy/io/fits/hdu/compressed.py
@@ -85,6 +85,9 @@ class CompImageHeader(Header):
     This essentially wraps the image header, so that all values are read from
     and written to the image header.  However, updates to the image header will
     also update the table header where appropriate.
+
+    Note that if no image header is passed in, the code will instantiate a
+    regular `~astropy.io.fits.Header`.
     """
 
     # TODO: The difficulty of implementing this screams a need to rewrite this
@@ -105,9 +108,19 @@ class CompImageHeader(Header):
     # the schema system, but it's not quite ready for that yet.  Also it still
     # makes more sense to change CompImageHDU to subclass ImageHDU :/
 
-    def __init__(self, table_header, image_header=None):
+    def __new__(cls, table_header, image_header=None):
+        # 2019-09-14 (MHvK): No point wrapping anything if no image_header is
+        # given.  This happens if __getitem__ and copy are called - our super
+        # class will aim to initialize a new, possibly partially filled
+        # header, but we cannot usefully deal with that.
+        # TODO: the above suggests strongly we should *not* subclass from
+        # Header.  See also comment above about the need for reorganization.
         if image_header is None:
-            image_header = Header()
+            return Header(table_header)
+        else:
+            return super().__new__(cls)
+
+    def __init__(self, table_header, image_header):
         self._cards = image_header._cards
         self._keyword_indices = image_header._keyword_indices
         self._rvkc_indices = image_header._rvkc_indices
diff --git a/astropy/io/fits/header.py b/astropy/io/fits/header.py
index abd253e4ad..f60ef93049 100644
--- a/astropy/io/fits/header.py
+++ b/astropy/io/fits/header.py
@@ -132,10 +132,10 @@ class Header:
 
     def __getitem__(self, key):
         if isinstance(key, slice):
-            return Header([copy.copy(c) for c in self._cards[key]])
+            return self.__class__([copy.copy(c) for c in self._cards[key]])
         elif self._haswildcard(key):
-            return Header([copy.copy(self._cards[idx])
-                           for idx in self._wildcardmatch(key)])
+            return self.__class__([copy.copy(self._cards[idx])
+                                   for idx in self._wildcardmatch(key)])
         elif (isinstance(key, str) and
               key.upper() in Card._commentary_keywords):
             key = key.upper()
@@ -824,7 +824,7 @@ class Header:
             A new :class:`Header` instance.
         """
 
-        tmp = Header(copy.copy(card) for card in self._cards)
+        tmp = self.__class__((copy.copy(card) for card in self._cards))
         if strip:
             tmp._strip()
         return tmp
@@ -1294,7 +1294,7 @@ class Header:
             new cards to the header.
         """
 
-        temp = Header(cards)
+        temp = self.__class__(cards)
         if strip:
             temp._strip()
 


Test patch:
diff --git a/astropy/io/fits/tests/test_header.py b/astropy/io/fits/tests/test_header.py
index 47629fb773..a6093f5040 100644
--- a/astropy/io/fits/tests/test_header.py
+++ b/astropy/io/fits/tests/test_header.py
@@ -2838,3 +2838,36 @@ class TestRecordValuedKeywordCards(FitsTestCase):
         with open(self.temp('mode.fits'), mode=mode) as ff:
             hdu = fits.ImageHDU(data=np.ones(5))
             hdu.writeto(ff)
+
+
+def test_subclass():
+    """Check that subclasses don't get ignored on slicing and copying."""
+    class MyHeader(fits.Header):
+        def append(self, card, *args, **kwargs):
+            if isinstance(card, tuple) and len(card) == 2:
+                # Just for our checks we add a comment if there is none.
+                card += ('no comment',)
+
+            return super().append(card, *args, **kwargs)
+
+    my_header = MyHeader((('a', 1., 'first'),
+                          ('b', 2., 'second'),
+                          ('c', 3.,)))
+
+    assert my_header.comments['a'] == 'first'
+    assert my_header.comments['b'] == 'second'
+    assert my_header.comments['c'] == 'no comment'
+
+    slice_ = my_header[1:]
+    assert type(slice_) is MyHeader
+    assert slice_.comments['b'] == 'second'
+    assert slice_.comments['c'] == 'no comment'
+    selection = my_header['c*']
+    assert type(selection) is MyHeader
+    assert selection.comments['c'] == 'no comment'
+    copy_ = my_header.copy()
+    assert type(copy_) is MyHeader
+    assert copy_.comments['b'] == 'second'
+    assert copy_.comments['c'] == 'no comment'
+    my_header.extend((('d', 4.),))
+    assert my_header.comments['d'] == 'no comment'