Skip to content

Commit 7d03aa4

Browse files
committed
Add zstd support to Analyze format images.
- Add zstandard compression for analyze format images via pyzstd. - Add zstd option under setup.cfg options.extra_requires.
1 parent 62aea04 commit 7d03aa4

File tree

4 files changed

+30
-1
lines changed

4 files changed

+30
-1
lines changed

nibabel/analyze.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -907,6 +907,11 @@ class AnalyzeImage(SpatialImage):
907907
files_types = (('image', '.img'), ('header', '.hdr'))
908908
valid_exts = ('.img', '.hdr')
909909
_compressed_suffixes = ('.gz', '.bz2')
910+
try: # If pyzstd installed., add .zst suffix
911+
import pyzstd
912+
_compressed_suffixes = (*_compressed_suffixes, '.zst')
913+
except ImportError:
914+
pass
910915

911916
makeable = True
912917
rw = True

nibabel/openers.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,11 @@
3939
IndexedGzipFile = gzip.GzipFile
4040
HAVE_INDEXED_GZIP = False
4141

42+
# Enable .zst support if pyzstd installed.
43+
try:
44+
from pyzstd import ZstdFile
45+
except ImportError:
46+
pass
4247

4348
def _gzip_open(filename, mode='rb', compresslevel=9, keep_open=False):
4449

@@ -77,13 +82,19 @@ class Opener(object):
7782
"""
7883
gz_def = (_gzip_open, ('mode', 'compresslevel', 'keep_open'))
7984
bz2_def = (BZ2File, ('mode', 'buffering', 'compresslevel'))
85+
zstd_def = (ZstdFile, ('mode', 'level_or_option'))
8086
compress_ext_map = {
8187
'.gz': gz_def,
8288
'.bz2': bz2_def,
89+
'.zst': zstd_def,
8390
None: (open, ('mode', 'buffering')) # default
8491
}
8592
#: default compression level when writing gz and bz2 files
8693
default_compresslevel = 1
94+
#: default option for zst files
95+
default_zst_compresslevel = 3
96+
default_level_or_option = {"rb": None, "r": None,
97+
"wb": default_zst_compresslevel, "w": default_zst_compresslevel}
8798
#: whether to ignore case looking for compression extensions
8899
compress_ext_icase = True
89100

@@ -100,10 +111,15 @@ def __init__(self, fileish, *args, **kwargs):
100111
full_kwargs.update(dict(zip(arg_names[:n_args], args)))
101112
# Set default mode
102113
if 'mode' not in full_kwargs:
103-
kwargs['mode'] = 'rb'
114+
mode = 'rb'
115+
kwargs['mode'] = mode
116+
else:
117+
mode = full_kwargs['mode']
104118
# Default compression level
105119
if 'compresslevel' in arg_names and 'compresslevel' not in kwargs:
106120
kwargs['compresslevel'] = self.default_compresslevel
121+
if 'level_or_option' in arg_names and 'level_or_option' not in kwargs:
122+
kwargs['level_or_option'] = self.default_level_or_option[mode]
107123
# Default keep_open hint
108124
if 'keep_open' in arg_names:
109125
kwargs.setdefault('keep_open', False)

nibabel/volumeutils.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,12 @@
4040
#: file-like classes known to hold compressed data
4141
COMPRESSED_FILE_LIKES = (gzip.GzipFile, BZ2File, IndexedGzipFile)
4242

43+
# Enable .zst support if pyzstd installed.
44+
try:
45+
from pyzstd import ZstdFile
46+
COMPRESSED_FILE_LIKES = (*COMPRESSED_FILE_LIKES, ZstdFile)
47+
except ImportError:
48+
pass
4349

4450
class Recoder(object):
4551
""" class to return canonical code(s) from code or aliases

setup.cfg

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ test =
6060
pytest !=5.3.4
6161
pytest-cov
6262
pytest-doctestplus
63+
zstd =
64+
pyzstd >= 0.14.3
6365
all =
6466
%(dicomfs)s
6567
%(dev)s

0 commit comments

Comments
 (0)