TEST: Adjust existing tests/benchmarks to import zstd from _compression,

rather than calling optional_package again. Use "zstd" instead of "pyzstd"

Author: pauldmccarthy

nibabel/benchmarks/bench_fileslice.py

@@ -19,11 +19,11 @@
 from ..optpkg import optional_package
 from ..rstutils import rst_table
 from ..tmpdirs import InTemporaryDirectory
+from .._compression import HAVE_ZSTD
 
 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):

nibabel/tests/test_analyze.py

@@ -38,11 +38,10 @@
     suppress_warnings,
 )
 from ..tmpdirs import InTemporaryDirectory
+from .._compression import HAVE_ZSTD
 from . import test_spatialimages as tsi
 from . import test_wrapstruct as tws
 
-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'

nibabel/tests/test_loadsave.py

@@ -22,9 +22,9 @@
 from ..optpkg import optional_package
 from ..testing import deprecated_to, expires
 from ..tmpdirs import InTemporaryDirectory
+from .._compression import HAVE_ZSTD
 
 _, have_scipy, _ = optional_package('scipy')
-_, have_pyzstd, _ = optional_package('pyzstd')
 
 import pytest
 from numpy.testing import assert_almost_equal, assert_array_equal
@@ -84,7 +84,7 @@ def test_load_empty_image():
 
 @pytest.mark.parametrize('extension', ['.gz', '.bz2', '.zst'])
 def test_load_bad_compressed_extension(tmp_path, extension):
-    if extension == '.zst' and not have_pyzstd:
+    if extension == '.zst' and not HAVE_ZSTD:
         pytest.skip()
     file_path = tmp_path / f'img.nii{extension}'
     file_path.write_bytes(b'bad')
@@ -94,7 +94,7 @@ def test_load_bad_compressed_extension(tmp_path, extension):
 
 @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:
+    if extension == '.zst' and not HAVE_ZSTD:
         pytest.skip()
     file_path = tmp_path / f'img.nii{extension}'
     with Opener(file_path, 'wb') as fobj:

nibabel/tests/test_minc1.py

@@ -24,8 +24,7 @@
 from ..tmpdirs import InTemporaryDirectory
 from . import test_spatialimages as tsi
 from .test_fileslice import slicer_samples
-
-pyzstd, HAVE_ZSTD, _ = optional_package('pyzstd')
+from .._compression import zstd, HAVE_ZSTD
 
 EG_FNAME = pjoin(data_path, 'tiny.mnc')
 
@@ -175,7 +174,7 @@ def test_compressed(self):
             content = open(tp['fname'], 'rb').read()
             openers_exts = [(gzip.open, '.gz'), (bz2.BZ2File, '.bz2')]
             if HAVE_ZSTD:  # add .zst to test if installed
-                openers_exts += [(pyzstd.ZstdFile, '.zst')]
+                openers_exts += [(zstd.ZstdFile, '.zst')]
             with InTemporaryDirectory():
                 for opener, ext in openers_exts:
                     fname = 'test.mnc' + ext

nibabel/tests/test_openers.py

@@ -23,8 +23,7 @@
 from ..openers import HAVE_INDEXED_GZIP, BZ2File, DeterministicGzipFile, ImageOpener, Opener
 from ..optpkg import optional_package
 from ..tmpdirs import InTemporaryDirectory
-
-pyzstd, HAVE_ZSTD, _ = optional_package('pyzstd')
+from .._compression import zstd, HAVE_ZSTD
 
 
 class Lunk:
@@ -271,7 +270,7 @@ class StrictOpener(Opener):
                     IndexedGzipFile = GzipFile
                 assert isinstance(fobj.fobj, (GzipFile, IndexedGzipFile))
             elif lext == 'zst':
-                assert isinstance(fobj.fobj, pyzstd.ZstdFile)
+                assert isinstance(fobj.fobj, zstd.ZstdFile)
             else:
                 assert isinstance(fobj.fobj, BZ2File)
 

nibabel/tests/test_volumeutils.py

@@ -56,8 +56,7 @@
     working_type,
     write_zeros,
 )
-
-pyzstd, HAVE_ZSTD, _ = optional_package('pyzstd')
+from .._compression import zstd, HAVE_ZSTD
 
 # convenience variables for numpy types
 FLOAT_TYPES = sctypes['float']
@@ -81,7 +80,7 @@ def test__is_compressed_fobj():
     with InTemporaryDirectory():
         file_openers = [('', open, False), ('.gz', gzip.open, True), ('.bz2', BZ2File, True)]
         if HAVE_ZSTD:
-            file_openers += [('.zst', pyzstd.ZstdFile, True)]
+            file_openers += [('.zst', zstd.ZstdFile, True)]
         for ext, opener, compressed in file_openers:
             fname = 'test.bin' + ext
             for mode in ('wb', 'rb'):
@@ -105,7 +104,7 @@ def make_array(n, bytes):
     with InTemporaryDirectory():
         openers = [open, gzip.open, BZ2File]
         if HAVE_ZSTD:
-            openers += [pyzstd.ZstdFile]
+            openers += [zstd.ZstdFile]
         for n, opener in itertools.product((256, 1024, 2560, 25600), openers):
             in_arr = np.arange(n, dtype=dtype)
             # Write array to file
@@ -262,7 +261,7 @@ def test_array_from_file_reread():
     with InTemporaryDirectory():
         openers = [open, gzip.open, bz2.BZ2File, BytesIO]
         if HAVE_ZSTD:
-            openers += [pyzstd.ZstdFile]
+            openers += [zstd.ZstdFile]
         for shape, opener, dtt, order in itertools.product(
             ((64,), (64, 65), (64, 65, 66)), openers, (np.int16, np.float32), ('F', 'C')
         ):