TYP: Annotate openers

Opener proxy methods now match io.BufferedIOBase prototypes.

Remove some version checks for indexed-gzip < 0.8, which supported
Python 3.6 while our minimum is now 3.8.

A runtime-checkable protocol for .read()/.write() was the easiest way to
accommodate weird file-likes that aren't IOBases.

When indexed-gzip is typed, we may need to adjust the output of
_gzip_open.

Author: effigies

nibabel/openers.py

@@ -7,34 +7,48 @@
 #
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Context manager openers for various fileobject types"""
+from __future__ import annotations
+
 import gzip
-import warnings
+import io
+import typing as ty
 from bz2 import BZ2File
 from os.path import splitext
 
-from packaging.version import Version
-
 from nibabel.optpkg import optional_package
 
-# is indexed_gzip present and modern?
-try:
-    import indexed_gzip as igzip  # type: ignore
+if ty.TYPE_CHECKING:  # pragma: no cover
+    from types import TracebackType
 
-    version = igzip.__version__
+    import pyzstd
+    from _typeshed import WriteableBuffer
 
-    HAVE_INDEXED_GZIP = True
+    ModeRT = ty.Literal['r', 'rt']
+    ModeRB = ty.Literal['rb']
+    ModeWT = ty.Literal['w', 'wt']
+    ModeWB = ty.Literal['wb']
+    ModeR = ty.Union[ModeRT, ModeRB]
+    ModeW = ty.Union[ModeWT, ModeWB]
+    Mode = ty.Union[ModeR, ModeW]
 
-    # < 0.7 - no good
-    if Version(version) < Version('0.7.0'):
-        warnings.warn(f'indexed_gzip is present, but too old (>= 0.7.0 required): {version})')
-        HAVE_INDEXED_GZIP = False
-    # >= 0.8 SafeIndexedGzipFile renamed to IndexedGzipFile
-    elif Version(version) < Version('0.8.0'):
-        IndexedGzipFile = igzip.SafeIndexedGzipFile
-    else:
-        IndexedGzipFile = igzip.IndexedGzipFile
-    del igzip, version
+    OpenerDef = tuple[ty.Callable[..., io.IOBase], tuple[str, ...]]
+else:
+    pyzstd = optional_package('pyzstd')[0]
+
+
+@ty.runtime_checkable
+class Fileish(ty.Protocol):
+    def read(self, size: int = -1, /) -> bytes:
+        ...  # pragma: no cover
+
+    def write(self, b: bytes, /) -> int | None:
+        ...  # pragma: no cover
+
+
+try:
+    from indexed_gzip import IndexedGzipFile  # type: ignore
 
+    HAVE_INDEXED_GZIP = True
 except ImportError:
     # nibabel.openers.IndexedGzipFile is imported by nibabel.volumeutils
     # to detect compressed file types, so we give a fallback value here.
@@ -49,35 +63,63 @@ class DeterministicGzipFile(gzip.GzipFile):
     to a modification time (``mtime``) of 0 seconds.
     """
 
-    def __init__(self, filename=None, mode=None, compresslevel=9, fileobj=None, mtime=0):
-        # These two guards are copied from
+    def __init__(
+        self,
+        filename: str | None = None,
+        mode: Mode | None = None,
+        compresslevel: int = 9,
+        fileobj: io.FileIO | None = None,
+        mtime: int = 0,
+    ):
+        if mode is None:
+            mode = 'rb'
+        modestr: str = mode
+
+        # These two guards are adapted from
         # https://github.com/python/cpython/blob/6ab65c6/Lib/gzip.py#L171-L174
-        if mode and 'b' not in mode:
-            mode += 'b'
+        if 'b' not in modestr:
+            modestr = f'{mode}b'
         if fileobj is None:
-            fileobj = self.myfileobj = open(filename, mode or 'rb')
+            if filename is None:
+                raise TypeError('Must define either fileobj or filename')
+            # Cast because GzipFile.myfileobj has type io.FileIO while open returns ty.IO
+            fileobj = self.myfileobj = ty.cast(io.FileIO, open(filename, modestr))
         return super().__init__(
-            filename='', mode=mode, compresslevel=compresslevel, fileobj=fileobj, mtime=mtime
+            filename='',
+            mode=modestr,
+            compresslevel=compresslevel,
+            fileobj=fileobj,
+            mtime=mtime,
         )
 
 
-def _gzip_open(filename, mode='rb', compresslevel=9, mtime=0, keep_open=False):
+def _gzip_open(
+    filename: str,
+    mode: Mode = 'rb',
+    compresslevel: int = 9,
+    mtime: int = 0,
+    keep_open: bool = False,
+) -> gzip.GzipFile:
+
+    if not HAVE_INDEXED_GZIP or mode != 'rb':
+        gzip_file = DeterministicGzipFile(filename, mode, compresslevel, mtime=mtime)
 
     # use indexed_gzip if possible for faster read access.  If keep_open ==
     # True, we tell IndexedGzipFile to keep the file handle open. Otherwise
     # the IndexedGzipFile will close/open the file on each read.
-    if HAVE_INDEXED_GZIP and mode == 'rb':
-        gzip_file = IndexedGzipFile(filename, drop_handles=not keep_open)
-
-    # Fall-back to built-in GzipFile
     else:
-        gzip_file = DeterministicGzipFile(filename, mode, compresslevel, mtime=mtime)
+        gzip_file = IndexedGzipFile(filename, drop_handles=not keep_open)
 
     return gzip_file
 
 
-def _zstd_open(filename, mode='r', *, level_or_option=None, zstd_dict=None):
-    pyzstd = optional_package('pyzstd')[0]
+def _zstd_open(
+    filename: str,
+    mode: Mode = 'r',
+    *,
+    level_or_option: int | dict | None = None,
+    zstd_dict: pyzstd.ZstdDict | None = None,
+) -> pyzstd.ZstdFile:
     return pyzstd.ZstdFile(filename, mode, level_or_option=level_or_option, zstd_dict=zstd_dict)
 
 
@@ -104,7 +146,7 @@ class Opener:
     gz_def = (_gzip_open, ('mode', 'compresslevel', 'mtime', 'keep_open'))
     bz2_def = (BZ2File, ('mode', 'buffering', 'compresslevel'))
     zstd_def = (_zstd_open, ('mode', 'level_or_option', 'zstd_dict'))
-    compress_ext_map = {
+    compress_ext_map: dict[str | None, OpenerDef] = {
         '.gz': gz_def,
         '.bz2': bz2_def,
         '.zst': zstd_def,
@@ -121,19 +163,19 @@ class Opener:
         'w': default_zst_compresslevel,
     }
     #: whether to ignore case looking for compression extensions
-    compress_ext_icase = True
+    compress_ext_icase: bool = True
+
+    fobj: io.IOBase
 
-    def __init__(self, fileish, *args, **kwargs):
-        if self._is_fileobj(fileish):
+    def __init__(self, fileish: str | io.IOBase, *args, **kwargs):
+        if isinstance(fileish, (io.IOBase, Fileish)):
             self.fobj = fileish
             self.me_opened = False
-            self._name = None
+            self._name = getattr(fileish, 'name', None)
             return
         opener, arg_names = self._get_opener_argnames(fileish)
         # Get full arguments to check for mode and compresslevel
-        full_kwargs = kwargs.copy()
-        n_args = len(args)
-        full_kwargs.update(dict(zip(arg_names[:n_args], args)))
+        full_kwargs = {**kwargs, **dict(zip(arg_names, args))}
         # Set default mode
         if 'mode' not in full_kwargs:
             mode = 'rb'
@@ -155,7 +197,7 @@ def __init__(self, fileish, *args, **kwargs):
         self._name = fileish
         self.me_opened = True
 
-    def _get_opener_argnames(self, fileish):
+    def _get_opener_argnames(self, fileish: str) -> OpenerDef:
         _, ext = splitext(fileish)
         if self.compress_ext_icase:
             ext = ext.lower()
@@ -168,16 +210,12 @@ def _get_opener_argnames(self, fileish):
             return self.compress_ext_map[ext]
         return self.compress_ext_map[None]
 
-    def _is_fileobj(self, obj):
-        """Is `obj` a file-like object?"""
-        return hasattr(obj, 'read') and hasattr(obj, 'write')
-
     @property
-    def closed(self):
+    def closed(self) -> bool:
         return self.fobj.closed
 
     @property
-    def name(self):
+    def name(self) -> str | None:
         """Return ``self.fobj.name`` or self._name if not present
 
         self._name will be None if object was created with a fileobj, otherwise
@@ -186,42 +224,53 @@ def name(self):
         return self._name
 
     @property
-    def mode(self):
-        return self.fobj.mode
+    def mode(self) -> str:
+        # Check and raise our own error for type narrowing purposes
+        if hasattr(self.fobj, 'mode'):
+            return self.fobj.mode
+        raise AttributeError(f'{self.fobj.__class__.__name__} has no attribute "mode"')
 
-    def fileno(self):
+    def fileno(self) -> int:
         return self.fobj.fileno()
 
-    def read(self, *args, **kwargs):
-        return self.fobj.read(*args, **kwargs)
+    def read(self, size: int = -1, /) -> bytes:
+        return self.fobj.read(size)
 
-    def readinto(self, *args, **kwargs):
-        return self.fobj.readinto(*args, **kwargs)
+    def readinto(self, buffer: WriteableBuffer, /) -> int | None:
+        # Check and raise our own error for type narrowing purposes
+        if hasattr(self.fobj, 'readinto'):
+            return self.fobj.readinto(buffer)
+        raise AttributeError(f'{self.fobj.__class__.__name__} has no attribute "readinto"')
 
-    def write(self, *args, **kwargs):
-        return self.fobj.write(*args, **kwargs)
+    def write(self, b: bytes, /) -> int | None:
+        return self.fobj.write(b)
 
-    def seek(self, *args, **kwargs):
-        return self.fobj.seek(*args, **kwargs)
+    def seek(self, pos: int, whence: int = 0, /) -> int:
+        return self.fobj.seek(pos, whence)
 
-    def tell(self, *args, **kwargs):
-        return self.fobj.tell(*args, **kwargs)
+    def tell(self, /) -> int:
+        return self.fobj.tell()
 
-    def close(self, *args, **kwargs):
-        return self.fobj.close(*args, **kwargs)
+    def close(self, /) -> None:
+        return self.fobj.close()
 
-    def __iter__(self):
+    def __iter__(self) -> ty.Iterator[bytes]:
         return iter(self.fobj)
 
-    def close_if_mine(self):
+    def close_if_mine(self) -> None:
         """Close ``self.fobj`` iff we opened it in the constructor"""
         if self.me_opened:
             self.close()
 
-    def __enter__(self):
+    def __enter__(self) -> Opener:
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         self.close_if_mine()
 
 

nibabel/tests/test_openers.py

@@ -38,7 +38,7 @@ def __init__(self, message):
     def write(self):
         pass
 
-    def read(self):
+    def read(self, size=-1, /):
         return self.message