TYP: Annotate Cifti2Extension

Author: effigies

nibabel/cifti2/cifti2.py

@@ -1570,7 +1570,7 @@ def to_file_map(self, file_map=None, dtype=None):
 
         self.update_headers()
         header = self._nifti_header
-        extension = Cifti2Extension(content=self.header.to_xml())
+        extension = Cifti2Extension.from_bytes(self.header.to_xml())
         header.extensions = Nifti1Extensions(
             ext for ext in header.extensions if not isinstance(ext, Cifti2Extension)
         )

nibabel/cifti2/parse_cifti2.py

@@ -40,19 +40,15 @@
 )
 
 
-class Cifti2Extension(Nifti1Extension):
+class Cifti2Extension(Nifti1Extension[Cifti2Header]):
     code = 32
 
-    def __init__(self, code=None, content=None):
-        Nifti1Extension.__init__(self, code=code or self.code, content=content)
-
-    def _unmangle(self, value):
+    def _unmangle(self, value: bytes) -> Cifti2Header:
         parser = Cifti2Parser()
         parser.parse(string=value)
-        self._content = parser.header
-        return self._content
+        return parser.header
 
-    def _mangle(self, value):
+    def _mangle(self, value: Cifti2Header) -> bytes:
         if not isinstance(value, Cifti2Header):
             raise ValueError('Can only mangle a Cifti2Header.')
         return value.to_xml()