RF: Add generic NiftiExtension base class

Nifti1Extension is a non-ideal base class for NIfTI extensions because
it assumes that it is safe to store use a null transformation, and thus
default to `bytes` objects. This makes it difficult to define its typing
behavior in a way that allows subclasses to refine the type such that
type-checkers understand it.

This patch creates a generic `NiftiExtension` class that parameterizes
the "runtime representation" type. Nifti1Extension subclasses with
another parameter that defaults to `bytes`, allowing it to be subclassed
in turn (preserving the Nifti1Extension -> Nifti1DicomExtension subclass
relationship) while still emitting `bytes`.

We could have simply made `Nifti1Extension` the base class, but the
mangle/unmangle methods need some casts or ignore comments to type-check
cleanly. This separation allows us to have a clean base class with the
legacy hacks cordoned off into an subclass.

Author: effigies

nibabel/nifti1.py

@@ -13,11 +13,13 @@
 
 from __future__ import annotations
 
+import typing as ty
 import warnings
 from io import BytesIO
 
 import numpy as np
 import numpy.linalg as npl
+from typing_extensions import TypeVar  # PY312
 
 from . import analyze  # module import
 from .arrayproxy import get_obj_dtype
@@ -31,7 +33,19 @@
 from .spm99analyze import SpmAnalyzeHeader
 from .volumeutils import Recoder, endian_codes, make_dt_codes
 
-pdcm, have_dicom, _ = optional_package('pydicom')
+if ty.TYPE_CHECKING:
+    import pydicom as pdcm
+
+    have_dicom = True
+    DicomDataset = pdcm.Dataset
+else:
+    pdcm, have_dicom, _ = optional_package('pydicom')
+    if have_dicom:
+        DicomDataset = pdcm.Dataset
+    else:
+        DicomDataset = ty.Any
+
+T = TypeVar('T', default=bytes)
 
 # nifti1 flat header definition for Analyze-like first 348 bytes
 # first number in comments indicates offset in file header in bytes
@@ -283,110 +297,103 @@
 )
 
 
-class Nifti1Extension:
-    """Baseclass for NIfTI1 header extensions.
+class NiftiExtension(ty.Generic[T]):
+    """Base class for NIfTI header extensions."""
 
-    This class is sufficient to handle very simple text-based extensions, such
-    as `comment`. More sophisticated extensions should/will be supported by
-    dedicated subclasses.
-    """
+    code: int
+    encoding: ty.Optional[str] = None
+    _content: bytes
+    _object: ty.Optional[T] = None
 
-    def __init__(self, code, content):
+    def __init__(
+        self,
+        code: ty.Union[int, str],
+        content: bytes,
+    ) -> None:
         """
         Parameters
         ----------
         code : int or str
           Canonical extension code as defined in the NIfTI standard, given
           either as integer or corresponding label
           (see :data:`~nibabel.nifti1.extension_codes`)
-        content : str
-          Extension content as read from the NIfTI file header. This content is
-          converted into a runtime representation.
+        content : bytes
+          Extension content as read from the NIfTI file header. This content may
+          be converted into a runtime representation.
         """
         try:
-            self._code = extension_codes.code[code]
+            self.code = extension_codes.code[code]  # type: ignore[assignment]
         except KeyError:
-            # XXX or fail or at least complain?
-            self._code = code
-        self._content = self._unmangle(content)
+            self.code = code  # type: ignore[assignment]
+        self._content = content
 
-    def _unmangle(self, value):
-        """Convert the extension content into its runtime representation.
+    # Handle (de)serialization of extension content
+    # Subclasses may implement these methods to provide an alternative
+    # view of the extension content. If left unimplemented, the content
+    # must be bytes and is not modified.
+    def _mangle(self, obj: T) -> bytes:
+        raise NotImplementedError
 
-        The default implementation does nothing at all.
+    def _unmangle(self, content: bytes) -> T:
+        raise NotImplementedError
 
-        Parameters
-        ----------
-        value : str
-          Extension content as read from file.
+    def _sync(self) -> None:
+        """Synchronize content with object.
 
-        Returns
-        -------
-        The same object that was passed as `value`.
-
-        Notes
-        -----
-        Subclasses should reimplement this method to provide the desired
-        unmangling procedure and may return any type of object.
+        This permits the runtime representation to be modified in-place
+        and updates the bytes representation accordingly.
         """
-        return value
-
-    def _mangle(self, value):
-        """Convert the extension content into NIfTI file header representation.
+        if self._object is not None:
+            self._content = self._mangle(self._object)
 
-        The default implementation does nothing at all.
-
-        Parameters
-        ----------
-        value : str
-          Extension content in runtime form.
+    def __repr__(self) -> str:
+        try:
+            code = extension_codes.label[self.code]
+        except KeyError:
+            # deal with unknown codes
+            code = self.code
+        return f'{self.__class__.__name__}({code}, {self._content!r})'
 
-        Returns
-        -------
-        str
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, self.__class__)
+            and self.code == other.code
+            and self.content == other.content
+        )
 
-        Notes
-        -----
-        Subclasses should reimplement this method to provide the desired
-        mangling procedure.
-        """
-        return value
+    def __ne__(self, other):
+        return not self == other
 
     def get_code(self):
         """Return the canonical extension type code."""
-        return self._code
+        return self.code
 
-    def get_content(self):
-        """Return the extension content in its runtime representation."""
+    @property
+    def content(self) -> bytes:
+        """Return the extension content as raw bytes."""
+        self._sync()
         return self._content
 
-    def get_sizeondisk(self):
+    def get_content(self) -> T:
+        """Return the extension content in its runtime representation.
+
+        This method may return a different type for each extension type.
+        """
+        if self._object is None:
+            self._object = self._unmangle(self._content)
+        return self._object
+
+    def get_sizeondisk(self) -> int:
         """Return the size of the extension in the NIfTI file."""
+        self._sync()
         # need raw value size plus 8 bytes for esize and ecode
-        size = len(self._mangle(self._content))
-        size += 8
+        size = len(self._content) + 8
         # extensions size has to be a multiple of 16 bytes
         if size % 16 != 0:
             size += 16 - (size % 16)
         return size
 
-    def __repr__(self):
-        try:
-            code = extension_codes.label[self._code]
-        except KeyError:
-            # deal with unknown codes
-            code = self._code
-
-        s = f"Nifti1Extension('{code}', '{self._content}')"
-        return s
-
-    def __eq__(self, other):
-        return (self._code, self._content) == (other._code, other._content)
-
-    def __ne__(self, other):
-        return not self == other
-
-    def write_to(self, fileobj, byteswap):
+    def write_to(self, fileobj: ty.BinaryIO, byteswap: bool = False) -> None:
         """Write header extensions to fileobj
 
         Write starts at fileobj current file position.
@@ -402,22 +409,74 @@ def write_to(self, fileobj, byteswap):
         -------
         None
         """
+        self._sync()
         extstart = fileobj.tell()
         rawsize = self.get_sizeondisk()
         # write esize and ecode first
-        extinfo = np.array((rawsize, self._code), dtype=np.int32)
+        extinfo = np.array((rawsize, self.code), dtype=np.int32)
         if byteswap:
             extinfo = extinfo.byteswap()
         fileobj.write(extinfo.tobytes())
         # followed by the actual extension content
         # XXX if mangling upon load is implemented, it should be reverted here
-        fileobj.write(self._mangle(self._content))
+        fileobj.write(self._content)
         # be nice and zero out remaining part of the extension till the
         # next 16 byte border
         fileobj.write(b'\x00' * (extstart + rawsize - fileobj.tell()))
 
 
-class Nifti1DicomExtension(Nifti1Extension):
+class Nifti1Extension(NiftiExtension[T]):
+    """Baseclass for NIfTI1 header extensions.
+
+    This class is sufficient to handle very simple text-based extensions, such
+    as `comment`. More sophisticated extensions should/will be supported by
+    dedicated subclasses.
+    """
+
+    def _unmangle(self, value: bytes) -> T:
+        """Convert the extension content into its runtime representation.
+
+        The default implementation does nothing at all.
+
+        Parameters
+        ----------
+        value : str
+          Extension content as read from file.
+
+        Returns
+        -------
+        The same object that was passed as `value`.
+
+        Notes
+        -----
+        Subclasses should reimplement this method to provide the desired
+        unmangling procedure and may return any type of object.
+        """
+        return value  # type: ignore[return-value]
+
+    def _mangle(self, value: T) -> bytes:
+        """Convert the extension content into NIfTI file header representation.
+
+        The default implementation does nothing at all.
+
+        Parameters
+        ----------
+        value : str
+          Extension content in runtime form.
+
+        Returns
+        -------
+        str
+
+        Notes
+        -----
+        Subclasses should reimplement this method to provide the desired
+        mangling procedure.
+        """
+        return value  # type: ignore[return-value]
+
+
+class Nifti1DicomExtension(Nifti1Extension[DicomDataset]):
     """NIfTI1 DICOM header extension
 
     This class is a thin wrapper around pydicom to read a binary DICOM
@@ -427,7 +486,12 @@ class Nifti1DicomExtension(Nifti1Extension):
     header.
     """
 
-    def __init__(self, code, content, parent_hdr=None):
+    def __init__(
+        self,
+        code: ty.Union[int, str],
+        content: ty.Union[bytes, DicomDataset, None] = None,
+        parent_hdr: ty.Optional[Nifti1Header] = None,
+    ) -> None:
         """
         Parameters
         ----------
@@ -452,50 +516,48 @@ def __init__(self, code, content, parent_hdr=None):
         code should always be 2 for DICOM.
         """
 
-        self._code = code
-        if parent_hdr:
-            self._is_little_endian = parent_hdr.endianness == '<'
-        else:
-            self._is_little_endian = True
+        self._is_little_endian = parent_hdr is None or parent_hdr.endianness == '<'
+
+        bytes_content: bytes
         if isinstance(content, pdcm.dataset.Dataset):
             self._is_implicit_VR = False
-            self._raw_content = self._mangle(content)
-            self._content = content
+            self._object = content
+            bytes_content = self._mangle(content)
         elif isinstance(content, bytes):  # Got a byte string - unmangle it
-            self._raw_content = content
-            self._is_implicit_VR = self._guess_implicit_VR()
-            ds = self._unmangle(content, self._is_implicit_VR, self._is_little_endian)
-            self._content = ds
+            self._is_implicit_VR = self._guess_implicit_VR(content)
+            self._object = self._unmangle(content)
+            bytes_content = content
         elif content is None:  # initialize a new dicom dataset
             self._is_implicit_VR = False
-            self._content = pdcm.dataset.Dataset()
+            self._object = pdcm.dataset.Dataset()
+            bytes_content = self._mangle(self._object)
         else:
             raise TypeError(
                 f'content must be either a bytestring or a pydicom Dataset. '
                 f'Got {content.__class__}'
             )
+        super().__init__(code, bytes_content)
 
-    def _guess_implicit_VR(self):
+    @staticmethod
+    def _guess_implicit_VR(content) -> bool:
         """Try to guess DICOM syntax by checking for valid VRs.
 
         Without a DICOM Transfer Syntax, it's difficult to tell if Value
         Representations (VRs) are included in the DICOM encoding or not.
         This reads where the first VR would be and checks it against a list of
         valid VRs
         """
-        potential_vr = self._raw_content[4:6].decode()
-        if potential_vr in pdcm.values.converters.keys():
-            implicit_VR = False
-        else:
-            implicit_VR = True
-        return implicit_VR
-
-    def _unmangle(self, value, is_implicit_VR=False, is_little_endian=True):
-        bio = BytesIO(value)
-        ds = pdcm.filereader.read_dataset(bio, is_implicit_VR, is_little_endian)
-        return ds
+        potential_vr = content[4:6].decode()
+        return potential_vr not in pdcm.values.converters.keys()
+
+    def _unmangle(self, obj: bytes) -> DicomDataset:
+        return pdcm.filereader.read_dataset(
+            BytesIO(obj),
+            self._is_implicit_VR,
+            self._is_little_endian,
+        )
 
-    def _mangle(self, dataset):
+    def _mangle(self, dataset: DicomDataset) -> bytes:
         bio = BytesIO()
         dio = pdcm.filebase.DicomFileLike(bio)
         dio.is_implicit_VR = self._is_implicit_VR