Add zstd to all; Add tests for zstd

Author: vanandrew

nibabel/tests/test_analyze.py

@@ -787,7 +787,7 @@ def test_big_offset_exts(self):
         arr = np.arange(24, dtype=np.int16).reshape((2, 3, 4))
         aff = np.eye(4)
         img_ext = img_klass.files_types[0][1]
-        compressed_exts = ['', '.gz', '.bz2']
+        compressed_exts = ['', '.gz', '.bz2', '.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 ZstdFile
 
 from ..tmpdirs import InTemporaryDirectory
 from ..deprecator import ExpiredDeprecationError
@@ -170,7 +171,9 @@ 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'),
+                            (ZstdFile, '.zst'))
             with InTemporaryDirectory():
                 for opener, ext in openers_exts:
                     fname = 'test.mnc' + ext

nibabel/tests/test_openers.py

@@ -74,6 +74,7 @@ def test_Opener_various():
         for input in ('test.txt',
                       'test.txt.gz',
                       'test.txt.bz2',
+                      'test.txt.zst',
                       sobj):
             with Opener(input, 'wb') as fobj:
                 fobj.write(message)
@@ -276,6 +277,7 @@ def test_name():
         for input in ('test.txt',
                       'test.txt.gz',
                       'test.txt.bz2',
+                      'test.txt.zst',
                       sobj,
                       lunk):
             exp_name = input if type(input) == type('') else None
@@ -332,6 +334,7 @@ def test_iter():
         for input, does_t in (('test.txt', True),
                               ('test.txt.gz', False),
                               ('test.txt.bz2', False),
+                              ('test.txt.zst', False),
                               (sobj, True)):
             with Opener(input, 'wb') as fobj:
                 for line in lines:

nibabel/tests/test_volumeutils.py

@@ -45,7 +45,7 @@
                            _write_data,
                            _ftype4scaled_finite,
                            )
-from ..openers import Opener, BZ2File
+from ..openers import Opener, BZ2File, ZstdFile
 from ..casting import (floor_log2, type_info, OK_FLOATS, shared_range)
 
 from ..deprecator import ExpiredDeprecationError
@@ -70,7 +70,8 @@ def test__is_compressed_fobj():
     with InTemporaryDirectory():
         for ext, opener, compressed in (('', open, False),
                                         ('.gz', gzip.open, True),
-                                        ('.bz2', BZ2File, True)):
+                                        ('.bz2', BZ2File, True),
+                                        ('.zst', ZstdFile, True)):
             fname = 'test.bin' + ext
             for mode in ('wb', 'rb'):
                 fobj = opener(fname, mode)
@@ -88,12 +89,12 @@ 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():
         for n, opener in itertools.product(
                 (256, 1024, 2560, 25600),
-                (open, gzip.open, BZ2File)):
+                (open, gzip.open, BZ2File, ZstdFile)):
             in_arr = np.arange(n, dtype=dtype)
             # Write array to file
             fobj_w = opener(fname, 'wb')
@@ -230,7 +231,7 @@ 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'),
+        for ext, offset in itertools.product(('', '.gz', '.bz2', '.zst'),
                                              (0, 5, 10)):
             fname = 'test.bin' + ext
             with Opener(fname, 'wb') as out_buf:
@@ -253,7 +254,7 @@ def test_array_from_file_reread():
     with InTemporaryDirectory():
         for shape, opener, dtt, order in itertools.product(
                 ((64,), (64, 65), (64, 65, 66)),
-                (open, gzip.open, bz2.BZ2File, BytesIO),
+                (open, gzip.open, bz2.BZ2File, ZstdFile, BytesIO),
                 (np.int16, np.float32),
                 ('F', 'C')):
             n_els = np.prod(shape)
@@ -901,7 +902,7 @@ 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', 'test.zst'
     start = 10
     end = 100
     diff = end - start
@@ -920,9 +921,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 +950,7 @@ 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'):
+        for in_file in ('test2.gz', 'test2.bz2', 'test2.zst'):
             # 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 =