Merge branch 'master' into fix/gzip_compression_options

Author: effigies

nibabel/analyze.py

@@ -906,7 +906,7 @@ class AnalyzeImage(SpatialImage):
     _meta_sniff_len = header_class.sizeof_hdr
     files_types = (('image', '.img'), ('header', '.hdr'))
     valid_exts = ('.img', '.hdr')
-    _compressed_suffixes = ('.gz', '.bz2')
+    _compressed_suffixes = ('.gz', '.bz2', '.zst')
 
     makeable = True
     rw = True

nibabel/benchmarks/bench_fileslice.py

@@ -18,13 +18,15 @@
 from ..fileslice import fileslice
 from ..rstutils import rst_table
 from ..tmpdirs import InTemporaryDirectory
+from ..optpkg import optional_package
 
 SHAPE = (64, 64, 32, 100)
 ROW_NAMES = [f'axis {i}, len {dim}' for i, dim in enumerate(SHAPE)]
 COL_NAMES = ['mid int',
              'step 1',
              'half step 1',
              'step mid int']
+HAVE_ZSTD = optional_package("pyzstd")[1]
 
 
 def _slices_for_len(L):
@@ -70,7 +72,8 @@ def g():
 def bench_fileslice(bytes=True,
                     file_=True,
                     gz=True,
-                    bz2=False):
+                    bz2=False,
+                    zst=True):
     sys.stdout.flush()
     repeat = 2
 
@@ -103,4 +106,10 @@ def my_table(title, times, base):
         my_table('bz2 slice - raw (ratio)',
                  np.dstack((bz2_times, bz2_times / bz2_base)),
                  bz2_base)
+    if zst and HAVE_ZSTD:
+        with InTemporaryDirectory():
+            zst_times, zst_base = run_slices('data.zst', repeat)
+        my_table('zst slice - raw (ratio)',
+                 np.dstack((zst_times, zst_times / zst_base)),
+                 zst_base)
     sys.stdout.flush()

nibabel/brikhead.py

@@ -490,7 +490,7 @@ class AFNIImage(SpatialImage):
     header_class = AFNIHeader
     valid_exts = ('.brik', '.head')
     files_types = (('image', '.brik'), ('header', '.head'))
-    _compressed_suffixes = ('.gz', '.bz2', '.Z')
+    _compressed_suffixes = ('.gz', '.bz2', '.Z', '.zst')
     makeable = False
     rw = False
     ImageArrayProxy = AFNIArrayProxy

nibabel/loadsave.py

@@ -19,6 +19,8 @@
 from .arrayproxy import is_proxy
 from .deprecated import deprecate_with_version
 
+_compressed_suffixes = ('.gz', '.bz2', '.zst')
+
 
 def _signature_matches_extension(filename, sniff):
     """Check if signature aka magic number matches filename extension.
@@ -153,7 +155,7 @@ def save(img, filename):
         return
 
     # Be nice to users by making common implicit conversions
-    froot, ext, trailing = splitext_addext(filename, ('.gz', '.bz2'))
+    froot, ext, trailing = splitext_addext(filename, _compressed_suffixes)
     lext = ext.lower()
 
     # Special-case Nifti singles and Pairs

nibabel/minc1.py

@@ -316,7 +316,7 @@ class Minc1Image(SpatialImage):
     _meta_sniff_len = 4
     valid_exts = ('.mnc',)
     files_types = (('image', '.mnc'),)
-    _compressed_suffixes = ('.gz', '.bz2')
+    _compressed_suffixes = ('.gz', '.bz2', '.zst')
 
     makeable = True
     rw = False

nibabel/openers.py

@@ -15,6 +15,8 @@
 from os.path import splitext
 from distutils.version import StrictVersion
 
+from nibabel.optpkg import optional_package
+
 # is indexed_gzip present and modern?
 try:
     import indexed_gzip as igzip
@@ -72,6 +74,12 @@ def _gzip_open(filename, mode='rb', compresslevel=9, mtime=0, keep_open=False):
     return gzip_file
 
 
+def _zstd_open(filename, mode="r", *, level_or_option=None, zstd_dict=None):
+    pyzstd = optional_package("pyzstd")[0]
+    return pyzstd.ZstdFile(filename, mode,
+                           level_or_option=level_or_option, zstd_dict=zstd_dict)
+
+
 class Opener(object):
     r""" Class to accept, maybe open, and context-manage file-likes / filenames
 
@@ -94,13 +102,20 @@ class Opener(object):
     """
     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 = {
         '.gz': gz_def,
         '.bz2': bz2_def,
+        '.zst': zstd_def,
         None: (open, ('mode', 'buffering'))  # default
     }
     #: default compression level when writing gz and bz2 files
     default_compresslevel = 1
+    #: default option for zst files
+    default_zst_compresslevel = 3
+    default_level_or_option = {"rb": None, "r": None,
+                               "wb": default_zst_compresslevel,
+                               "w": default_zst_compresslevel}
     #: whether to ignore case looking for compression extensions
     compress_ext_icase = True
 
@@ -117,10 +132,15 @@ def __init__(self, fileish, *args, **kwargs):
         full_kwargs.update(dict(zip(arg_names[:n_args], args)))
         # Set default mode
         if 'mode' not in full_kwargs:
-            kwargs['mode'] = 'rb'
+            mode = 'rb'
+            kwargs['mode'] = mode
+        else:
+            mode = full_kwargs['mode']
         # Default compression level
         if 'compresslevel' in arg_names and 'compresslevel' not in kwargs:
             kwargs['compresslevel'] = self.default_compresslevel
+        if 'level_or_option' in arg_names and 'level_or_option' not in kwargs:
+            kwargs['level_or_option'] = self.default_level_or_option[mode]
         # Default keep_open hint
         if 'keep_open' in arg_names:
             kwargs.setdefault('keep_open', False)

nibabel/tests/test_analyze.py

@@ -30,6 +30,7 @@
 from ..casting import as_int
 from ..tmpdirs import InTemporaryDirectory
 from ..arraywriters import WriterError
+from ..optpkg import optional_package
 
 import pytest
 from numpy.testing import (assert_array_equal, assert_array_almost_equal)
@@ -40,6 +41,8 @@
 from .test_wrapstruct import _TestLabeledWrapStruct
 from . import test_spatialimages as tsi
 
+HAVE_ZSTD = optional_package("pyzstd")[1]
+
 header_file = os.path.join(data_path, 'analyze.hdr')
 
 PIXDIM0_MSG = 'pixdim[1,2,3] should be non-zero; setting 0 dims to 1'
@@ -788,6 +791,8 @@ def test_big_offset_exts(self):
         aff = np.eye(4)
         img_ext = img_klass.files_types[0][1]
         compressed_exts = ['', '.gz', '.bz2']
+        if HAVE_ZSTD:
+            compressed_exts += ['.zst']
         with InTemporaryDirectory():
             for offset in (0, 2048):
                 # Set offset in in-memory image

nibabel/tests/test_minc1.py

@@ -22,6 +22,7 @@
 from ..deprecated import ModuleProxy
 from .. import minc1
 from ..minc1 import Minc1File, Minc1Image, MincHeader
+from ..optpkg import optional_package
 
 from ..tmpdirs import InTemporaryDirectory
 from ..deprecator import ExpiredDeprecationError
@@ -32,6 +33,8 @@
 from . import test_spatialimages as tsi
 from .test_fileslice import slicer_samples
 
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
+
 EG_FNAME = pjoin(data_path, 'tiny.mnc')
 
 # Example images in format expected for ``test_image_api``, adding ``zooms``
@@ -170,7 +173,10 @@ def test_compressed(self):
         # Not so for MINC2; hence this small sub-class
         for tp in self.test_files:
             content = open(tp['fname'], 'rb').read()
-            openers_exts = ((gzip.open, '.gz'), (bz2.BZ2File, '.bz2'))
+            openers_exts = [(gzip.open, '.gz'),
+                            (bz2.BZ2File, '.bz2')]
+            if HAVE_ZSTD:  # add .zst to test if installed
+                openers_exts += [(pyzstd.ZstdFile, '.zst')]
             with InTemporaryDirectory():
                 for opener, ext in openers_exts:
                     fname = 'test.mnc' + ext

nibabel/tests/test_openers.py

@@ -16,15 +16,23 @@
 import time
 
 from numpy.compat.py3k import asstr, asbytes
-from ..openers import Opener, ImageOpener, HAVE_INDEXED_GZIP, BZ2File, DeterministicGzipFile
+from ..openers import (Opener,
+                       ImageOpener,
+                       HAVE_INDEXED_GZIP,
+                       BZ2File,
+                       DeterministicGzipFile,
+                       )
 from ..tmpdirs import InTemporaryDirectory
 from ..volumeutils import BinOpener
+from ..optpkg import optional_package
 
 import unittest
 from unittest import mock
 import pytest
 from ..testing import error_warnings
 
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
+
 
 class Lunk(object):
     # bare file-like for testing
@@ -73,10 +81,13 @@ def test_Opener_various():
         import indexed_gzip as igzip
     with InTemporaryDirectory():
         sobj = BytesIO()
-        for input in ('test.txt',
-                      'test.txt.gz',
-                      'test.txt.bz2',
-                      sobj):
+        files_to_test = ['test.txt',
+                         'test.txt.gz',
+                         'test.txt.bz2',
+                         sobj]
+        if HAVE_ZSTD:
+            files_to_test += ['test.txt.zst']
+        for input in files_to_test:
             with Opener(input, 'wb') as fobj:
                 fobj.write(message)
                 assert fobj.tell() == len(message)
@@ -242,6 +253,8 @@ def test_compressed_ext_case():
     class StrictOpener(Opener):
         compress_ext_icase = False
     exts = ('gz', 'bz2', 'GZ', 'gZ', 'BZ2', 'Bz2')
+    if HAVE_ZSTD:
+        exts += ('zst', 'ZST', 'Zst')
     with InTemporaryDirectory():
         # Make a basic file to check type later
         with open(__file__, 'rb') as a_file:
@@ -266,6 +279,8 @@ class StrictOpener(Opener):
                 except ImportError:
                     IndexedGzipFile = GzipFile
                 assert isinstance(fobj.fobj, (GzipFile, IndexedGzipFile))
+            elif lext == 'zst':
+                assert isinstance(fobj.fobj, pyzstd.ZstdFile)
             else:
                 assert isinstance(fobj.fobj, BZ2File)
 
@@ -275,11 +290,14 @@ def test_name():
     sobj = BytesIO()
     lunk = Lunk('in ART')
     with InTemporaryDirectory():
-        for input in ('test.txt',
-                      'test.txt.gz',
-                      'test.txt.bz2',
-                      sobj,
-                      lunk):
+        files_to_test = ['test.txt',
+                         'test.txt.gz',
+                         'test.txt.bz2',
+                         sobj,
+                         lunk]
+        if HAVE_ZSTD:
+            files_to_test += ['test.txt.zst']
+        for input in files_to_test:
             exp_name = input if type(input) == type('') else None
             with Opener(input, 'wb') as fobj:
                 assert fobj.name == exp_name
@@ -331,10 +349,13 @@ def test_iter():
 """.split('\n')
     with InTemporaryDirectory():
         sobj = BytesIO()
-        for input, does_t in (('test.txt', True),
-                              ('test.txt.gz', False),
-                              ('test.txt.bz2', False),
-                              (sobj, True)):
+        files_to_test = [('test.txt', True),
+                         ('test.txt.gz', False),
+                         ('test.txt.bz2', False),
+                         (sobj, True)]
+        if HAVE_ZSTD:
+            files_to_test += [('test.txt.zst', False)]
+        for input, does_t in files_to_test:
             with Opener(input, 'wb') as fobj:
                 for line in lines:
                     fobj.write(asbytes(line + os.linesep))

nibabel/tests/test_volumeutils.py

@@ -49,13 +49,16 @@
 from ..casting import (floor_log2, type_info, OK_FLOATS, shared_range)
 
 from ..deprecator import ExpiredDeprecationError
+from ..optpkg import optional_package
 
 from numpy.testing import (assert_array_almost_equal,
                            assert_array_equal)
 import pytest
 
 from nibabel.testing import nullcontext, assert_dt_equal, assert_allclose_safely, suppress_warnings
 
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
+
 #: convenience variables for numpy types
 FLOAT_TYPES = np.sctypes['float']
 COMPLEX_TYPES = np.sctypes['complex']
@@ -68,9 +71,12 @@
 def test__is_compressed_fobj():
     # _is_compressed helper function
     with InTemporaryDirectory():
-        for ext, opener, compressed in (('', open, False),
-                                        ('.gz', gzip.open, True),
-                                        ('.bz2', BZ2File, True)):
+        file_openers = [('', open, False),
+                        ('.gz', gzip.open, True),
+                        ('.bz2', BZ2File, True)]
+        if HAVE_ZSTD:
+            file_openers += [('.zst', pyzstd.ZstdFile, True)]
+        for ext, opener, compressed in file_openers:
             fname = 'test.bin' + ext
             for mode in ('wb', 'rb'):
                 fobj = opener(fname, mode)
@@ -88,12 +94,15 @@ def make_array(n, bytes):
         arr.flags.writeable = True
         return arr
 
-    # Check whether file, gzip file, bz2 file reread memory from cache
+    # Check whether file, gzip file, bz2, zst file reread memory from cache
     fname = 'test.bin'
     with InTemporaryDirectory():
+        openers = [open, gzip.open, BZ2File]
+        if HAVE_ZSTD:
+            openers += [pyzstd.ZstdFile]
         for n, opener in itertools.product(
                 (256, 1024, 2560, 25600),
-                (open, gzip.open, BZ2File)):
+                openers):
             in_arr = np.arange(n, dtype=dtype)
             # Write array to file
             fobj_w = opener(fname, 'wb')
@@ -230,7 +239,10 @@ def test_array_from_file_openers():
     dtype = np.dtype(np.float32)
     in_arr = np.arange(24, dtype=dtype).reshape(shape)
     with InTemporaryDirectory():
-        for ext, offset in itertools.product(('', '.gz', '.bz2'),
+        extensions = ['', '.gz', '.bz2']
+        if HAVE_ZSTD:
+            extensions += ['.zst']
+        for ext, offset in itertools.product(extensions,
                                              (0, 5, 10)):
             fname = 'test.bin' + ext
             with Opener(fname, 'wb') as out_buf:
@@ -251,9 +263,12 @@ def test_array_from_file_reread():
     offset = 9
     fname = 'test.bin'
     with InTemporaryDirectory():
+        openers = [open, gzip.open, bz2.BZ2File, BytesIO]
+        if HAVE_ZSTD:
+            openers += [pyzstd.ZstdFile]
         for shape, opener, dtt, order in itertools.product(
                 ((64,), (64, 65), (64, 65, 66)),
-                (open, gzip.open, bz2.BZ2File, BytesIO),
+                openers,
                 (np.int16, np.float32),
                 ('F', 'C')):
             n_els = np.prod(shape)
@@ -901,7 +916,9 @@ def test_write_zeros():
 def test_seek_tell():
     # Test seek tell routine
     bio = BytesIO()
-    in_files = bio, 'test.bin', 'test.gz', 'test.bz2'
+    in_files = [bio, 'test.bin', 'test.gz', 'test.bz2']
+    if HAVE_ZSTD:
+        in_files += ['test.zst']
     start = 10
     end = 100
     diff = end - start
@@ -920,9 +937,12 @@ def test_seek_tell():
                 fobj.write(b'\x01' * start)
                 assert fobj.tell() == start
                 # Files other than BZ2Files can seek forward on write, leaving
-                # zeros in their wake.  BZ2Files can't seek when writing, unless
-                # we enable the write0 flag to seek_tell
-                if not write0 and in_file == 'test.bz2':  # Can't seek write in bz2
+                # zeros in their wake.  BZ2Files can't seek when writing,
+                # unless we enable the write0 flag to seek_tell
+                # ZstdFiles also does not support seek forward on write
+                if (not write0 and
+                    (in_file == 'test.bz2' or
+                     in_file == 'test.zst')):  # Can't seek write in bz2, zst
                     # write the zeros by hand for the read test below
                     fobj.write(b'\x00' * diff)
                 else:
@@ -946,7 +966,10 @@ def test_seek_tell():
             # Check we have the expected written output
             with ImageOpener(in_file, 'rb') as fobj:
                 assert fobj.read() == b'\x01' * start + b'\x00' * diff + b'\x02' * tail
-        for in_file in ('test2.gz', 'test2.bz2'):
+        input_files = ['test2.gz', 'test2.bz2']
+        if HAVE_ZSTD:
+            input_files += ['test2.zst']
+        for in_file in input_files:
             # Check failure of write seek backwards
             with ImageOpener(in_file, 'wb') as fobj:
                 fobj.write(b'g' * 10)

nibabel/volumeutils.py

@@ -22,6 +22,9 @@
 from .openers import Opener, BZ2File, IndexedGzipFile
 from .deprecated import deprecate_with_version
 from .externals.oset import OrderedSet
+from .optpkg import optional_package
+
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
 
 sys_is_le = sys.byteorder == 'little'
 native_code = sys_is_le and '<' or '>'
@@ -40,6 +43,10 @@
 #: file-like classes known to hold compressed data
 COMPRESSED_FILE_LIKES = (gzip.GzipFile, BZ2File, IndexedGzipFile)
 
+# Enable .zst support if pyzstd installed.
+if HAVE_ZSTD:
+    COMPRESSED_FILE_LIKES = (*COMPRESSED_FILE_LIKES, pyzstd.ZstdFile)
+
 
 class Recoder(object):
     """ class to return canonical code(s) from code or aliases

setup.cfg

@@ -61,6 +61,8 @@ test =
     pytest !=5.3.4
     pytest-cov
     pytest-doctestplus
+zstd =
+    pyzstd >= 0.14.3
 all =
     %(dicomfs)s
     %(dev)s
@@ -69,6 +71,7 @@ all =
     %(spm)s
     %(style)s
     %(test)s
+    %(zstd)s
 
 [options.entry_points]
 console_scripts =

tools/ci/env.sh

@@ -5,7 +5,7 @@ REQUIREMENTS="-r requirements.txt"
 # Minimum versions of minimum requirements
 MIN_REQUIREMENTS="-r min-requirements.txt"
 
-DEFAULT_OPT_DEPENDS="scipy matplotlib pillow pydicom h5py indexed_gzip"
+DEFAULT_OPT_DEPENDS="scipy matplotlib pillow pydicom h5py indexed_gzip pyzstd"
 # pydicom has skipped some important pre-releases, so enable a check against master
 PYDICOM_MASTER="git+https://github.com/pydicom/pydicom.git@master"
 # Minimum versions of optional requirements