Merge pull request #1 from effigies/zstd_support

RF: Use optional_package to allow code to assume pyzstd is present

Author: vanandrew

nibabel/analyze.py

@@ -95,7 +95,6 @@
 from .fileholders import copy_file_map
 from .batteryrunners import Report
 from .arrayproxy import ArrayProxy
-from .openers import HAVE_ZSTD
 
 # Sub-parts of standard analyze header from
 # Mayo dbh.h file
@@ -907,9 +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')
-    if HAVE_ZSTD:  # If pyzstd installed., add .zst suffix
-        _compressed_suffixes = (*_compressed_suffixes, '.zst')
+    _compressed_suffixes = ('.gz', '.bz2', '.zst')
 
     makeable = True
     rw = True

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/brikhead.py

@@ -43,7 +43,6 @@
     ImageDataError
 )
 from .volumeutils import Recoder
-from .openers import HAVE_ZSTD
 
 # used for doc-tests
 filepath = os.path.dirname(os.path.realpath(__file__))
@@ -491,9 +490,7 @@ class AFNIImage(SpatialImage):
     header_class = AFNIHeader
     valid_exts = ('.brik', '.head')
     files_types = (('image', '.brik'), ('header', '.head'))
-    _compressed_suffixes = ('.gz', '.bz2', '.Z')
-    if HAVE_ZSTD:  # If pyzstd installed., add .zst suffix
-        _compressed_suffixes = (*_compressed_suffixes, '.zst')
+    _compressed_suffixes = ('.gz', '.bz2', '.Z', '.zst')
     makeable = False
     rw = False
     ImageArrayProxy = AFNIArrayProxy

nibabel/loadsave.py

@@ -13,15 +13,13 @@
 import numpy as np
 
 from .filename_parser import splitext_addext, _stringify_path
-from .openers import ImageOpener, HAVE_ZSTD
+from .openers import ImageOpener
 from .filebasedimages import ImageFileError
 from .imageclasses import all_image_classes
 from .arrayproxy import is_proxy
 from .deprecated import deprecate_with_version
 
-_compressed_suffixes = ('.gz', '.bz2')
-if HAVE_ZSTD:  # If pyzstd installed., add .zst suffix
-    _compressed_suffixes = (*_compressed_suffixes, '.zst')
+_compressed_suffixes = ('.gz', '.bz2', '.zst')
 
 
 def load(filename, **kwargs):

nibabel/minc1.py

@@ -16,7 +16,6 @@
 
 from .spatialimages import SpatialHeader, SpatialImage
 from .fileslice import canonical_slicers
-from .openers import HAVE_ZSTD
 
 _dt_dict = {
     ('b', 'unsigned'): np.uint8,
@@ -317,9 +316,7 @@ class Minc1Image(SpatialImage):
     _meta_sniff_len = 4
     valid_exts = ('.mnc',)
     files_types = (('image', '.mnc'),)
-    _compressed_suffixes = ('.gz', '.bz2')
-    if HAVE_ZSTD:  # If pyzstd installed., add .zst suffix
-        _compressed_suffixes = (*_compressed_suffixes, '.zst')
+    _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
@@ -39,13 +41,6 @@
     IndexedGzipFile = gzip.GzipFile
     HAVE_INDEXED_GZIP = False
 
-# Enable .zst support if pyzstd installed.
-try:
-    from pyzstd import ZstdFile
-    HAVE_ZSTD = True
-except ImportError:
-    HAVE_ZSTD = False
-
 
 def _gzip_open(filename, mode='rb', compresslevel=9, keep_open=False):
 
@@ -62,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
 
@@ -84,14 +85,13 @@ 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
     }
-    if HAVE_ZSTD:  # add zst to ext map, if library exists
-        zstd_def = (ZstdFile, ('mode', 'level_or_option'))
-        compress_ext_map['.zst'] = zstd_def
     #: default compression level when writing gz and bz2 files
     default_compresslevel = 1
     #: default option for zst files

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,