Merge pull request #1138 from effigies/fix/1137

FIX: Testing for extension/signature mismatches shouldn't use sniff

Author: effigies

nibabel/filebasedimages.py

@@ -430,7 +430,7 @@ def _sniff_meta_for(klass, filename, sniff_nbytes, sniff=None):
         try:
             with ImageOpener(meta_fname, 'rb') as fobj:
                 binaryblock = fobj.read(sniff_nbytes)
-        except IOError:
+        except (IOError, EOFError):
             return None
         return (binaryblock, meta_fname)
 

nibabel/filename_parser.py

@@ -253,7 +253,7 @@ def _iendswith(whole, end):
 
 
 def splitext_addext(filename,
-                    addexts=('.gz', '.bz2'),
+                    addexts=('.gz', '.bz2', '.zst'),
                     match_case=False):
     """ Split ``/pth/fname.ext.gz`` into ``/pth/fname, .ext, .gz``
 

nibabel/loadsave.py

@@ -22,18 +22,14 @@
 _compressed_suffixes = ('.gz', '.bz2', '.zst')
 
 
-def _signature_matches_extension(filename, sniff):
+def _signature_matches_extension(filename):
     """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
@@ -48,20 +44,20 @@ def _signature_matches_extension(filename, sniff):
     """
     signatures = {
         ".gz": {"signature": b"\x1f\x8b", "format_name": "gzip"},
-        ".bz2": {"signature": b"BZh", "format_name": "bzip2"}
+        ".bz2": {"signature": b"BZh", "format_name": "bzip2"},
+        ".zst": {"signature": b"\x28\xb5\x2f\xfd", "format_name": "ztsd"},
     }
     filename = _stringify_path(filename)
     *_, ext = splitext_addext(filename)
     ext = ext.lower()
     if ext not in signatures:
         return True, ""
     expected_signature = signatures[ext]["signature"]
-    if sniff is None or len(sniff) < len(expected_signature):
-        try:
-            with open(filename, "rb") as fh:
-                sniff = fh.read(len(expected_signature))
-        except OSError:
-            return False, f"Could not read file: {filename}"
+    try:
+        with open(filename, "rb") as fh:
+            sniff = fh.read(len(expected_signature))
+    except OSError:
+        return False, f"Could not read file: {filename}"
     if sniff.startswith(expected_signature):
         return True, ""
     format_name = signatures[ext]["format_name"]
@@ -100,7 +96,7 @@ def load(filename, **kwargs):
             img = image_klass.from_filename(filename, **kwargs)
             return img
 
-    matches, msg = _signature_matches_extension(filename, sniff)
+    matches, msg = _signature_matches_extension(filename)
     if not matches:
         raise ImageFileError(msg)
 

nibabel/tests/test_loadsave.py

@@ -13,9 +13,11 @@
 from ..loadsave import load, read_img_data, _signature_matches_extension
 from ..filebasedimages import ImageFileError
 from ..tmpdirs import InTemporaryDirectory, TemporaryDirectory
+from ..openers import Opener
 
 from ..optpkg import optional_package
 _, have_scipy, _ = optional_package('scipy')
+_, have_pyzstd, _ = optional_package('pyzstd')
 
 from numpy.testing import (assert_almost_equal,
                            assert_array_equal)
@@ -74,41 +76,46 @@ def test_load_empty_image():
     assert str(err.value).startswith('Empty file: ')
 
 
-@pytest.mark.parametrize("extension", [".gz", ".bz2"])
+@pytest.mark.parametrize("extension", [".gz", ".bz2", ".zst"])
 def test_load_bad_compressed_extension(tmp_path, extension):
+    if extension == ".zst" and not have_pyzstd:
+        pytest.skip()
     file_path = tmp_path / f"img.nii{extension}"
     file_path.write_bytes(b"bad")
     with pytest.raises(ImageFileError, match=".*is not a .* file"):
         load(file_path)
 
 
+@pytest.mark.parametrize("extension", [".gz", ".bz2", ".zst"])
+def test_load_good_extension_with_bad_data(tmp_path, extension):
+    if extension == ".zst" and not have_pyzstd:
+        pytest.skip()
+    file_path = tmp_path / f"img.nii{extension}"
+    with Opener(file_path, "wb") as fobj:
+        fobj.write(b"bad")
+    with pytest.raises(ImageFileError, match="Cannot work out file type of .*"):
+        load(file_path)
+
+
 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)
+    matches, msg = _signature_matches_extension(tmp_path / "uncompressed.nii")
     assert matches
     assert msg == ""
-    matches, msg = _signature_matches_extension(tmp_path / "missing.gz", None)
+    matches, msg = _signature_matches_extension(tmp_path / "missing.gz")
     assert not matches
     assert msg.startswith("Could not read")
-    matches, msg = _signature_matches_extension(bad_file, None)
+    matches, msg = _signature_matches_extension(bad_file)
     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])
+    matches, msg = _signature_matches_extension(good_file)
     assert matches
     assert msg == ""
-    matches, msg = _signature_matches_extension(good_file, None)
+    matches, msg = _signature_matches_extension(tmp_path / "missing.nii")
     assert matches
     assert msg == ""