RF: Use optional_package in tests/benchmark

Author: effigies

nibabel/benchmarks/bench_fileslice.py

@@ -14,17 +14,19 @@
 import numpy as np
 
 from io import BytesIO
-from ..openers import ImageOpener, HAVE_ZSTD
+from ..openers import ImageOpener
 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):
@@ -104,11 +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 HAVE_ZSTD:
-        if zst:
-            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)
+    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/tests/test_analyze.py

@@ -30,7 +30,7 @@
 from ..casting import as_int
 from ..tmpdirs import InTemporaryDirectory
 from ..arraywriters import WriterError
-from ..openers import HAVE_ZSTD
+from ..optpkg import optional_package
 
 import pytest
 from numpy.testing import (assert_array_equal, assert_array_almost_equal)
@@ -41,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'

nibabel/tests/test_minc1.py

@@ -22,7 +22,7 @@
 from ..deprecated import ModuleProxy
 from .. import minc1
 from ..minc1 import Minc1File, Minc1Image, MincHeader
-from ..openers import HAVE_ZSTD
+from ..optpkg import optional_package
 
 from ..tmpdirs import InTemporaryDirectory
 from ..deprecator import ExpiredDeprecationError
@@ -33,9 +33,7 @@
 from . import test_spatialimages as tsi
 from .test_fileslice import slicer_samples
 
-# only import ZstdFile, if installed
-if HAVE_ZSTD:
-    from ..openers import ZstdFile
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
 
 EG_FNAME = pjoin(data_path, 'tiny.mnc')
 
@@ -178,7 +176,7 @@ def test_compressed(self):
             openers_exts = [(gzip.open, '.gz'),
                             (bz2.BZ2File, '.bz2')]
             if HAVE_ZSTD:  # add .zst to test if installed
-                openers_exts += [(ZstdFile, '.zst')]
+                openers_exts += [(pyzstd.ZstdFile, '.zst')]
             with InTemporaryDirectory():
                 for opener, ext in openers_exts:
                     fname = 'test.mnc' + ext

nibabel/tests/test_openers.py

@@ -18,17 +18,17 @@
                        ImageOpener,
                        HAVE_INDEXED_GZIP,
                        BZ2File,
-                       HAVE_ZSTD)
+                       )
 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
 
-if HAVE_ZSTD:
-    from ..openers import ZstdFile
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
 
 
 class Lunk(object):
@@ -277,7 +277,7 @@ class StrictOpener(Opener):
                     IndexedGzipFile = GzipFile
                 assert isinstance(fobj.fobj, (GzipFile, IndexedGzipFile))
             elif lext == 'zst':
-                assert isinstance(fobj.fobj, ZstdFile)
+                assert isinstance(fobj.fobj, pyzstd.ZstdFile)
             else:
                 assert isinstance(fobj.fobj, BZ2File)
 

nibabel/tests/test_volumeutils.py

@@ -45,20 +45,19 @@
                            _write_data,
                            _ftype4scaled_finite,
                            )
-from ..openers import Opener, BZ2File, HAVE_ZSTD
+from ..openers import Opener, BZ2File
 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
 
-# only import ZstdFile, if installed
-if HAVE_ZSTD:
-    from ..openers import ZstdFile
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
 
 #: convenience variables for numpy types
 FLOAT_TYPES = np.sctypes['float']
@@ -76,7 +75,7 @@ def test__is_compressed_fobj():
                         ('.gz', gzip.open, True),
                         ('.bz2', BZ2File, True)]
         if HAVE_ZSTD:
-            file_openers += [('.zst', ZstdFile, True)]
+            file_openers += [('.zst', pyzstd.ZstdFile, True)]
         for ext, opener, compressed in file_openers:
             fname = 'test.bin' + ext
             for mode in ('wb', 'rb'):
@@ -100,7 +99,7 @@ def make_array(n, bytes):
     with InTemporaryDirectory():
         openers = [open, gzip.open, BZ2File]
         if HAVE_ZSTD:
-            openers += [ZstdFile]
+            openers += [pyzstd.ZstdFile]
         for n, opener in itertools.product(
                 (256, 1024, 2560, 25600),
                 openers):
@@ -266,7 +265,7 @@ def test_array_from_file_reread():
     with InTemporaryDirectory():
         openers = [open, gzip.open, bz2.BZ2File, BytesIO]
         if HAVE_ZSTD:
-            openers += [ZstdFile]
+            openers += [pyzstd.ZstdFile]
         for shape, opener, dtt, order in itertools.product(
                 ((64,), (64, 65), (64, 65, 66)),
                 openers,