Merge pull request #1005 from vanandrew/zstd_support

NF: Add zstd compression support

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
@@ -55,6 +57,12 @@ def _gzip_open(filename, mode='rb', compresslevel=9, 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
 
@@ -77,13 +85,20 @@ class Opener(object):
     """
     gz_def = (_gzip_open, ('mode', 'compresslevel', '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
 
@@ -100,10 +115,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

@@ -14,15 +14,22 @@
 from distutils.version import StrictVersion
 
 from numpy.compat.py3k import asstr, asbytes
-from ..openers import Opener, ImageOpener, HAVE_INDEXED_GZIP, BZ2File
+from ..openers import (Opener,
+                       ImageOpener,
+                       HAVE_INDEXED_GZIP,
+                       BZ2File,
+                       )
 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
@@ -71,10 +78,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)
@@ -240,6 +250,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:
@@ -264,6 +276,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)
 
@@ -273,11 +287,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
@@ -329,10 +346,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))