avoid extra read from image file

Author: jeromedockes

nibabel/loadsave.py

@@ -20,14 +20,18 @@
 from .deprecated import deprecate_with_version
 
 
-def _signature_matches_extension(filename):
+def _signature_matches_extension(filename, sniff):
     """Check if signature aka magic number matches filename extension.
 
     Parameters
     ----------
     filename : str or os.PathLike
         Path to the file to check
 
+    sniff : bytes or None
+        First bytes of the file. If not `None` and long enough to contain the
+        signature, avoids having to read the start of the file.
+
     Returns
     -------
     matches : bool
@@ -50,11 +54,14 @@ def _signature_matches_extension(filename):
     if ext not in signatures:
         return True, ""
     expected_signature = signatures[ext]["signature"]
-    try:
-        with open(filename, "rb") as fh:
-            found_signature = fh.read(len(expected_signature))
-    except OSError:
-        return False, f"Could not read file: {filename}"
+    if sniff is not None and len(sniff) >= len(expected_signature):
+        found_signature = sniff[:len(expected_signature)]
+    else:
+        try:
+            with open(filename, "rb") as fh:
+                found_signature = fh.read(len(expected_signature))
+        except OSError:
+            return False, f"Could not read file: {filename}"
     if found_signature == expected_signature:
         return True, ""
     format_name = signatures[ext]["format_name"]
@@ -85,9 +92,6 @@ def load(filename, **kwargs):
         raise FileNotFoundError(f"No such file or no access: '{filename}'")
     if stat_result.st_size <= 0:
         raise ImageFileError(f"Empty file: '{filename}'")
-    matches, msg = _signature_matches_extension(filename)
-    if not matches:
-        raise ImageFileError(msg)
 
     sniff = None
     for image_klass in all_image_classes:
@@ -96,6 +100,10 @@ def load(filename, **kwargs):
             img = image_klass.from_filename(filename, **kwargs)
             return img
 
+    matches, msg = _signature_matches_extension(filename, sniff)
+    if not matches:
+        raise ImageFileError(msg)
+
     raise ImageFileError(f'Cannot work out file type of "{filename}"')
 
 

nibabel/tests/test_loadsave.py

@@ -10,7 +10,7 @@
 from .. import (Spm99AnalyzeImage, Spm2AnalyzeImage,
                 Nifti1Pair, Nifti1Image,
                 Nifti2Pair, Nifti2Image)
-from ..loadsave import load, read_img_data
+from ..loadsave import load, read_img_data, _signature_matches_extension
 from ..filebasedimages import ImageFileError
 from ..tmpdirs import InTemporaryDirectory, TemporaryDirectory
 
@@ -84,11 +84,35 @@ def test_load_bad_compressed_extension(tmp_path, extension):
         load(file_path)
 
 
-def test_load_file_that_cannot_be_read(tmp_path):
-    subdir = tmp_path / "img.nii.gz"
-    subdir.mkdir()
-    with pytest.raises(ImageFileError, match="Could not read"):
-        load(subdir)
+def test_signature_matches_extension(tmp_path):
+    gz_signature = b"\x1f\x8b"
+    good_file = tmp_path / "good.gz"
+    good_file.write_bytes(gz_signature)
+    bad_file = tmp_path / "bad.gz"
+    bad_file.write_bytes(b"bad")
+    matches, msg = _signature_matches_extension(
+        tmp_path / "uncompressed.nii", None)
+    assert matches
+    assert msg == ""
+    matches, msg = _signature_matches_extension(tmp_path / "missing.gz", None)
+    assert not matches
+    assert msg.startswith("Could not read")
+    matches, msg = _signature_matches_extension(bad_file, None)
+    assert not matches
+    assert "is not a" in msg
+    matches, msg = _signature_matches_extension(bad_file, gz_signature + b"abc")
+    assert matches
+    assert msg == ""
+    matches, msg = _signature_matches_extension(
+        good_file, gz_signature + b"abc")
+    assert matches
+    assert msg == ""
+    matches, msg = _signature_matches_extension(good_file, gz_signature[:1])
+    assert matches
+    assert msg == ""
+    matches, msg = _signature_matches_extension(good_file, None)
+    assert matches
+    assert msg == ""
 
 
 def test_read_img_data_nifti():