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

Author: effigies

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

@@ -19,9 +19,12 @@
 import numpy as np
 
 from .casting import shared_range, OK_FLOATS
-from .openers import Opener, BZ2File, IndexedGzipFile, HAVE_ZSTD
+from .openers import Opener, BZ2File, IndexedGzipFile
 from .deprecated import deprecate_with_version
 from .externals.oset import OrderedSet
+from .optpkg import optional_package
+
+pyzstd, HAVE_ZSTD, _ = optional_package("pyzstd")
 
 sys_is_le = sys.byteorder == 'little'
 native_code = sys_is_le and '<' or '>'
@@ -42,8 +45,7 @@
 
 # Enable .zst support if pyzstd installed.
 if HAVE_ZSTD:
-    from .openers import ZstdFile
-    COMPRESSED_FILE_LIKES = (*COMPRESSED_FILE_LIKES, ZstdFile)
+    COMPRESSED_FILE_LIKES = (*COMPRESSED_FILE_LIKES, pyzstd.ZstdFile)
 
 
 class Recoder(object):