Merge pull request #1178 from effigies/type/primary_image_api

ENH: Add type annotations for image API (FileBasedImage and DataobjImage)

Author: effigies

nibabel/arrayproxy.py

@@ -25,6 +25,9 @@
 
 See :mod:`nibabel.tests.test_proxy_api` for proxy API conformance checks.
 """
+from __future__ import annotations
+
+import typing as ty
 import warnings
 from contextlib import contextmanager
 from threading import RLock
@@ -53,7 +56,28 @@
 KEEP_FILE_OPEN_DEFAULT = False
 
 
-class ArrayProxy:
+if ty.TYPE_CHECKING:  # pragma: no cover
+    import numpy.typing as npt
+
+
+class ArrayLike(ty.Protocol):
+    """Protocol for numpy ndarray-like objects
+
+    This is more stringent than :class:`numpy.typing.ArrayLike`, but guarantees
+    access to shape, ndim and slicing.
+    """
+
+    shape: tuple[int, ...]
+    ndim: int
+
+    def __array__(self, dtype: npt.DTypeLike | None = None, /) -> npt.NDArray:
+        ...  # pragma: no cover
+
+    def __getitem__(self, key, /) -> npt.NDArray:
+        ...  # pragma: no cover
+
+
+class ArrayProxy(ArrayLike):
     """Class to act as proxy for the array that can be read from a file
 
     The array proxy allows us to freeze the passed fileobj and header such that

nibabel/dataobj_images.py

@@ -7,17 +7,33 @@
 * returns an array from ``numpy.asanyarray(obj)``;
 * has an attribute or property ``shape``.
 """
+from __future__ import annotations
+
+import typing as ty
 
 import numpy as np
 
+from .arrayproxy import ArrayLike
 from .deprecated import deprecate_with_version
-from .filebasedimages import FileBasedImage
+from .filebasedimages import FileBasedHeader, FileBasedImage, FileMap, FileSpec
+
+if ty.TYPE_CHECKING:  # pragma: no cover
+    import numpy.typing as npt
 
 
 class DataobjImage(FileBasedImage):
     """Template class for images that have dataobj data stores"""
 
-    def __init__(self, dataobj, header=None, extra=None, file_map=None):
+    _data_cache: np.ndarray | None
+    _fdata_cache: np.ndarray | None
+
+    def __init__(
+        self,
+        dataobj: ArrayLike,
+        header: FileBasedHeader | ty.Mapping | None = None,
+        extra: ty.Mapping | None = None,
+        file_map: FileMap | None = None,
+    ):
         """Initialize dataobj image
 
         The datobj image is a combination of (dataobj, header), with optional
@@ -40,11 +56,11 @@ def __init__(self, dataobj, header=None, extra=None, file_map=None):
         """
         super().__init__(header=header, extra=extra, file_map=file_map)
         self._dataobj = dataobj
-        self._fdata_cache = None
         self._data_cache = None
+        self._fdata_cache = None
 
     @property
-    def dataobj(self):
+    def dataobj(self) -> ArrayLike:
         return self._dataobj
 
     @deprecate_with_version(
@@ -202,7 +218,11 @@ def get_data(self, caching='fill'):
             self._data_cache = data
         return data
 
-    def get_fdata(self, caching='fill', dtype=np.float64):
+    def get_fdata(
+        self,
+        caching: ty.Literal['fill', 'unchanged'] = 'fill',
+        dtype: npt.DTypeLike = np.float64,
+    ) -> np.ndarray:
         """Return floating point image data with necessary scaling applied
 
         The image ``dataobj`` property can be an array proxy or an array.  An
@@ -351,7 +371,7 @@ def get_fdata(self, caching='fill', dtype=np.float64):
         return data
 
     @property
-    def in_memory(self):
+    def in_memory(self) -> bool:
         """True when any array data is in memory cache
 
         There are separate caches for `get_data` reads and `get_fdata` reads.
@@ -363,7 +383,7 @@ def in_memory(self):
             or self._data_cache is not None
         )
 
-    def uncache(self):
+    def uncache(self) -> None:
         """Delete any cached read of data from proxied data
 
         Remember there are two types of images:
@@ -392,15 +412,21 @@ def uncache(self):
         self._data_cache = None
 
     @property
-    def shape(self):
+    def shape(self) -> tuple[int, ...]:
         return self._dataobj.shape
 
     @property
-    def ndim(self):
+    def ndim(self) -> int:
         return self._dataobj.ndim
 
     @classmethod
-    def from_file_map(klass, file_map, *, mmap=True, keep_file_open=None):
+    def from_file_map(
+        klass,
+        file_map: FileMap,
+        *,
+        mmap: bool | ty.Literal['c', 'r'] = True,
+        keep_file_open: bool | None = None,
+    ):
         """Class method to create image from mapping in ``file_map``
 
         Parameters
@@ -433,7 +459,13 @@ def from_file_map(klass, file_map, *, mmap=True, keep_file_open=None):
         raise NotImplementedError
 
     @classmethod
-    def from_filename(klass, filename, *, mmap=True, keep_file_open=None):
+    def from_filename(
+        klass,
+        filename: FileSpec,
+        *,
+        mmap: bool | ty.Literal['c', 'r'] = True,
+        keep_file_open: bool | None = None,
+    ):
         """Class method to create image from filename `filename`
 
         Parameters

nibabel/filebasedimages.py

@@ -10,6 +10,8 @@
 from __future__ import annotations
 
 import io
+import os
+import typing as ty
 from copy import deepcopy
 from typing import Type
 from urllib import request
@@ -18,6 +20,10 @@
 from .filename_parser import TypesFilenamesError, splitext_addext, types_filenames
 from .openers import ImageOpener
 
+FileSpec = ty.Union[str, os.PathLike]
+FileMap = ty.Mapping[str, FileHolder]
+FileSniff = ty.Tuple[bytes, str]
+
 
 class ImageFileError(Exception):
     pass
@@ -41,10 +47,10 @@ def from_header(klass, header=None):
         )
 
     @classmethod
-    def from_fileobj(klass, fileobj):
+    def from_fileobj(klass, fileobj: io.IOBase):
         raise NotImplementedError
 
-    def write_to(self, fileobj):
+    def write_to(self, fileobj: io.IOBase):
         raise NotImplementedError
 
     def __eq__(self, other):
@@ -53,7 +59,7 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
-    def copy(self):
+    def copy(self) -> FileBasedHeader:
         """Copy object to independent representation
 
         The copy should not be affected by any changes to the original
@@ -155,7 +161,12 @@ class FileBasedImage:
     makeable: bool = True  # Used in test code
     rw: bool = True  # Used in test code
 
-    def __init__(self, header=None, extra=None, file_map=None):
+    def __init__(
+        self,
+        header: FileBasedHeader | ty.Mapping | None = None,
+        extra: ty.Mapping | None = None,
+        file_map: FileMap | None = None,
+    ):
         """Initialize image
 
         The image is a combination of (header), with
@@ -182,14 +193,14 @@ def __init__(self, header=None, extra=None, file_map=None):
         self.file_map = file_map
 
     @property
-    def header(self):
+    def header(self) -> FileBasedHeader:
         return self._header
 
     def __getitem__(self, key):
         """No slicing or dictionary interface for images"""
         raise TypeError('Cannot slice image objects.')
 
-    def get_filename(self):
+    def get_filename(self) -> str | None:
         """Fetch the image filename
 
         Parameters
@@ -210,7 +221,7 @@ def get_filename(self):
         characteristic_type = self.files_types[0][0]
         return self.file_map[characteristic_type].filename
 
-    def set_filename(self, filename):
+    def set_filename(self, filename: str):
         """Sets the files in the object from a given filename
 
         The different image formats may check whether the filename has
@@ -228,16 +239,16 @@ def set_filename(self, filename):
         self.file_map = self.__class__.filespec_to_file_map(filename)
 
     @classmethod
-    def from_filename(klass, filename):
+    def from_filename(klass, filename: FileSpec):
         file_map = klass.filespec_to_file_map(filename)
         return klass.from_file_map(file_map)
 
     @classmethod
-    def from_file_map(klass, file_map):
+    def from_file_map(klass, file_map: FileMap):
         raise NotImplementedError
 
     @classmethod
-    def filespec_to_file_map(klass, filespec):
+    def filespec_to_file_map(klass, filespec: FileSpec):
         """Make `file_map` for this class from filename `filespec`
 
         Class method
@@ -271,7 +282,7 @@ def filespec_to_file_map(klass, filespec):
             file_map[key] = FileHolder(filename=fname)
         return file_map
 
-    def to_filename(self, filename, **kwargs):
+    def to_filename(self, filename: FileSpec, **kwargs):
         r"""Write image to files implied by filename string
 
         Parameters
@@ -290,11 +301,11 @@ def to_filename(self, filename, **kwargs):
         self.file_map = self.filespec_to_file_map(filename)
         self.to_file_map(**kwargs)
 
-    def to_file_map(self, file_map=None, **kwargs):
+    def to_file_map(self, file_map: FileMap | None = None, **kwargs):
         raise NotImplementedError
 
     @classmethod
-    def make_file_map(klass, mapping=None):
+    def make_file_map(klass, mapping: ty.Mapping[str, str | io.IOBase] | None = None):
         """Class method to make files holder for this image type
 
         Parameters
@@ -327,7 +338,7 @@ def make_file_map(klass, mapping=None):
     load = from_filename
 
     @classmethod
-    def instance_to_filename(klass, img, filename):
+    def instance_to_filename(klass, img: FileBasedImage, filename: FileSpec):
         """Save `img` in our own format, to name implied by `filename`
 
         This is a class method
@@ -343,7 +354,7 @@ def instance_to_filename(klass, img, filename):
         img.to_filename(filename)
 
     @classmethod
-    def from_image(klass, img):
+    def from_image(klass, img: FileBasedImage):
         """Class method to create new instance of own class from `img`
 
         Parameters
@@ -359,7 +370,12 @@ def from_image(klass, img):
         raise NotImplementedError()
 
     @classmethod
-    def _sniff_meta_for(klass, filename, sniff_nbytes, sniff=None):
+    def _sniff_meta_for(
+        klass,
+        filename: FileSpec,
+        sniff_nbytes: int,
+        sniff: FileSniff | None = None,
+    ):
         """Sniff metadata for image represented by `filename`
 
         Parameters
@@ -404,7 +420,12 @@ def _sniff_meta_for(klass, filename, sniff_nbytes, sniff=None):
         return (binaryblock, meta_fname)
 
     @classmethod
-    def path_maybe_image(klass, filename, sniff=None, sniff_max=1024):
+    def path_maybe_image(
+        klass,
+        filename: FileSpec,
+        sniff: FileSniff | None = None,
+        sniff_max: int = 1024,
+    ):
         """Return True if `filename` may be image matching this class
 
         Parameters
@@ -547,7 +568,7 @@ def from_bytes(klass, bytestring: bytes):
 
         Parameters
         ----------
-        bstring : bytes
+        bytestring : bytes
             Byte string containing the on-disk representation of an image
         """
         return klass.from_stream(io.BytesIO(bytestring))
@@ -571,7 +592,7 @@ def to_bytes(self, **kwargs) -> bytes:
         return bio.getvalue()
 
     @classmethod
-    def from_url(klass, url, timeout=5):
+    def from_url(klass, url: str | request.Request, timeout: float = 5):
         """Retrieve and load an image from a URL
 
         Class method