Add zstd to all; Add tests for zstd

Author: vanandrew

nibabel/tests/test_analyze.py

@@ -30,6 +30,7 @@
 from ..casting import as_int
 from ..tmpdirs import InTemporaryDirectory
 from ..arraywriters import WriterError
+from ..openers import HAVE_ZSTD
 
 import pytest
 from numpy.testing import (assert_array_equal, assert_array_almost_equal)
@@ -788,6 +789,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 ..openers import HAVE_ZSTD
 
 from ..tmpdirs import InTemporaryDirectory
 from ..deprecator import ExpiredDeprecationError
@@ -32,6 +33,10 @@
 from . import test_spatialimages as tsi
 from .test_fileslice import slicer_samples
 
+# only import ZstdFile, if installed
+if HAVE_ZSTD:
+    from ..openers import ZstdFile
+
 EG_FNAME = pjoin(data_path, 'tiny.mnc')
 
 # Example images in format expected for ``test_image_api``, adding ``zooms``
@@ -170,7 +175,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 += [(ZstdFile, '.zst')]
             with InTemporaryDirectory():
                 for opener, ext in openers_exts:
                     fname = 'test.mnc' + ext

nibabel/tests/test_openers.py

@@ -14,7 +14,11 @@
 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,
+                       HAVE_ZSTD)
 from ..tmpdirs import InTemporaryDirectory
 from ..volumeutils import BinOpener
 
@@ -23,6 +27,9 @@
 import pytest
 from ..testing import error_warnings
 
+if HAVE_ZSTD:
+    from ..openers import ZstdFile
+
 
 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, 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))

nibabel/tests/test_volumeutils.py

@@ -45,7 +45,7 @@
                            _write_data,
                            _ftype4scaled_finite,
                            )
-from ..openers import Opener, BZ2File
+from ..openers import Opener, BZ2File, HAVE_ZSTD
 from ..casting import (floor_log2, type_info, OK_FLOATS, shared_range)
 
 from ..deprecator import ExpiredDeprecationError
@@ -56,6 +56,10 @@
 
 from nibabel.testing import nullcontext, assert_dt_equal, assert_allclose_safely, suppress_warnings
 
+# only import ZstdFile, if installed
+if HAVE_ZSTD:
+    from ..openers import ZstdFile
+
 #: convenience variables for numpy types
 FLOAT_TYPES = np.sctypes['float']
 COMPLEX_TYPES = np.sctypes['complex']
@@ -68,9 +72,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', ZstdFile, True)]
+        for ext, opener, compressed in file_openers:
             fname = 'test.bin' + ext
             for mode in ('wb', 'rb'):
                 fobj = opener(fname, mode)
@@ -88,12 +95,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 += [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 +240,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 +264,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 += [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 +917,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 +938,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 +967,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)

setup.cfg

@@ -70,6 +70,7 @@ all =
     %(spm)s
     %(style)s
     %(test)s
+    %(zstd)s
 
 [options.entry_points]
 console_scripts =