Skip to content

Commit c2a5d4b

Browse files
AA-TurnerAA-Turner
authored
gh-132983: Clean-ups for _zstd (#133670)
1 parent a2a0fa9 commit c2a5d4b

File tree

5 files changed

+107
-176
lines changed

5 files changed

+107
-176
lines changed

‎Lib/compression/zstd/__init__.py

Lines changed: 40 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,16 @@
2828

2929
import _zstd
3030
import enum
31-
from _zstd import *
31+
from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdDict, ZstdError,
32+
get_frame_size, zstd_version)
3233
from compression.zstd._zstdfile import ZstdFile, open, _nbytes
3334

34-
COMPRESSION_LEVEL_DEFAULT = _zstd._compressionLevel_values[0]
35+
# zstd_version_number is (MAJOR * 100 * 100 + MINOR * 100 + RELEASE)
36+
zstd_version_info = (*divmod(_zstd.zstd_version_number // 100, 100),
37+
_zstd.zstd_version_number % 100)
38+
"""Version number of the runtime zstd library as a tuple of integers."""
39+
40+
COMPRESSION_LEVEL_DEFAULT = _zstd.ZSTD_CLEVEL_DEFAULT
3541
"""The default compression level for Zstandard, currently '3'."""
3642

3743

@@ -167,28 +173,28 @@ def decompress(data, zstd_dict=None, options=None):
167173
class CompressionParameter(enum.IntEnum):
168174
"""Compression parameters."""
169175

170-
compression_level = _zstd._ZSTD_c_compressionLevel
171-
window_log = _zstd._ZSTD_c_windowLog
172-
hash_log = _zstd._ZSTD_c_hashLog
173-
chain_log = _zstd._ZSTD_c_chainLog
174-
search_log = _zstd._ZSTD_c_searchLog
175-
min_match = _zstd._ZSTD_c_minMatch
176-
target_length = _zstd._ZSTD_c_targetLength
177-
strategy = _zstd._ZSTD_c_strategy
178-
179-
enable_long_distance_matching = _zstd._ZSTD_c_enableLongDistanceMatching
180-
ldm_hash_log = _zstd._ZSTD_c_ldmHashLog
181-
ldm_min_match = _zstd._ZSTD_c_ldmMinMatch
182-
ldm_bucket_size_log = _zstd._ZSTD_c_ldmBucketSizeLog
183-
ldm_hash_rate_log = _zstd._ZSTD_c_ldmHashRateLog
184-
185-
content_size_flag = _zstd._ZSTD_c_contentSizeFlag
186-
checksum_flag = _zstd._ZSTD_c_checksumFlag
187-
dict_id_flag = _zstd._ZSTD_c_dictIDFlag
188-
189-
nb_workers = _zstd._ZSTD_c_nbWorkers
190-
job_size = _zstd._ZSTD_c_jobSize
191-
overlap_log = _zstd._ZSTD_c_overlapLog
176+
compression_level = _zstd.ZSTD_c_compressionLevel
177+
window_log = _zstd.ZSTD_c_windowLog
178+
hash_log = _zstd.ZSTD_c_hashLog
179+
chain_log = _zstd.ZSTD_c_chainLog
180+
search_log = _zstd.ZSTD_c_searchLog
181+
min_match = _zstd.ZSTD_c_minMatch
182+
target_length = _zstd.ZSTD_c_targetLength
183+
strategy = _zstd.ZSTD_c_strategy
184+
185+
enable_long_distance_matching = _zstd.ZSTD_c_enableLongDistanceMatching
186+
ldm_hash_log = _zstd.ZSTD_c_ldmHashLog
187+
ldm_min_match = _zstd.ZSTD_c_ldmMinMatch
188+
ldm_bucket_size_log = _zstd.ZSTD_c_ldmBucketSizeLog
189+
ldm_hash_rate_log = _zstd.ZSTD_c_ldmHashRateLog
190+
191+
content_size_flag = _zstd.ZSTD_c_contentSizeFlag
192+
checksum_flag = _zstd.ZSTD_c_checksumFlag
193+
dict_id_flag = _zstd.ZSTD_c_dictIDFlag
194+
195+
nb_workers = _zstd.ZSTD_c_nbWorkers
196+
job_size = _zstd.ZSTD_c_jobSize
197+
overlap_log = _zstd.ZSTD_c_overlapLog
192198

193199
def bounds(self):
194200
"""Return the (lower, upper) int bounds of a compression parameter.
@@ -201,7 +207,7 @@ def bounds(self):
201207
class DecompressionParameter(enum.IntEnum):
202208
"""Decompression parameters."""
203209

204-
window_log_max = _zstd._ZSTD_d_windowLogMax
210+
window_log_max = _zstd.ZSTD_d_windowLogMax
205211

206212
def bounds(self):
207213
"""Return the (lower, upper) int bounds of a decompression parameter.
@@ -219,15 +225,15 @@ class Strategy(enum.IntEnum):
219225
the numeric value might change.
220226
"""
221227

222-
fast = _zstd._ZSTD_fast
223-
dfast = _zstd._ZSTD_dfast
224-
greedy = _zstd._ZSTD_greedy
225-
lazy = _zstd._ZSTD_lazy
226-
lazy2 = _zstd._ZSTD_lazy2
227-
btlazy2 = _zstd._ZSTD_btlazy2
228-
btopt = _zstd._ZSTD_btopt
229-
btultra = _zstd._ZSTD_btultra
230-
btultra2 = _zstd._ZSTD_btultra2
228+
fast = _zstd.ZSTD_fast
229+
dfast = _zstd.ZSTD_dfast
230+
greedy = _zstd.ZSTD_greedy
231+
lazy = _zstd.ZSTD_lazy
232+
lazy2 = _zstd.ZSTD_lazy2
233+
btlazy2 = _zstd.ZSTD_btlazy2
234+
btopt = _zstd.ZSTD_btopt
235+
btultra = _zstd.ZSTD_btultra
236+
btultra2 = _zstd.ZSTD_btultra2
231237

232238

233239
# Check validity of the CompressionParameter & DecompressionParameter types

‎Lib/compression/zstd/_zstdfile.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
import io
22
from os import PathLike
3-
from _zstd import (ZstdCompressor, ZstdDecompressor, _ZSTD_DStreamSizes,
4-
ZstdError)
3+
from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdError,
4+
ZSTD_DStreamOutSize)
55
from compression._common import _streams
66

77
__all__ = ("ZstdFile", "open")
88

9-
_ZSTD_DStreamOutSize = _ZSTD_DStreamSizes[1]
10-
119
_MODE_CLOSED = 0
1210
_MODE_READ = 1
1311
_MODE_WRITE = 2
@@ -188,7 +186,7 @@ def read1(self, size=-1):
188186
# Note this should *not* be io.DEFAULT_BUFFER_SIZE.
189187
# ZSTD_DStreamOutSize is the minimum amount to read guaranteeing
190188
# a full block is read.
191-
size = _ZSTD_DStreamOutSize
189+
size = ZSTD_DStreamOutSize
192190
return self._buffer.read1(size)
193191

194192
def readinto(self, b):

‎Lib/test/test_zstd.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,7 @@ def test_compress_parameters(self):
281281
with self.assertRaisesRegex(ZstdError,
282282
(r'Error when setting zstd compression parameter "window_log", '
283283
r'it should \d+ <= value <= \d+, provided value is 100\. '
284-
r'\(zstd v\d\.\d\.\d, (?:32|64)-bit build\)')):
284+
r'\((?:32|64)-bit build\)')):
285285
compress(b'', options=option)
286286

287287
def test_unknown_compression_parameter(self):
@@ -413,7 +413,7 @@ def test_decompress_parameters(self):
413413
with self.assertRaisesRegex(ZstdError,
414414
(r'Error when setting zstd decompression parameter "window_log_max", '
415415
r'it should \d+ <= value <= \d+, provided value is 100\. '
416-
r'\(zstd v\d\.\d\.\d, (?:32|64)-bit build\)')):
416+
r'\((?:32|64)-bit build\)')):
417417
decompress(b'', options=options)
418418

419419
def test_unknown_decompression_parameter(self):

‎Modules/_zstd/_zstdmodule.c

Lines changed: 61 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,6 @@ set_zstd_error(const _zstd_state* const state,
3333
case ERR_COMPRESS:
3434
msg = "Unable to compress zstd data: %s";
3535
break;
36-
case ERR_SET_PLEDGED_INPUT_SIZE:
37-
msg = "Unable to set pledged uncompressed content size: %s";
38-
break;
3936

4037
case ERR_LOAD_D_DICT:
4138
msg = "Unable to load zstd dictionary or prefix for decompression: %s";
@@ -151,19 +148,19 @@ set_parameter_error(const _zstd_state* const state, int is_compress,
151148
}
152149
if (ZSTD_isError(bounds.error)) {
153150
PyErr_Format(state->ZstdError,
154-
"Zstd %s parameter \"%s\" is invalid. (zstd v%s)",
155-
type, name, ZSTD_versionString());
151+
"Zstd %s parameter \"%s\" is invalid.",
152+
type, name);
156153
return;
157154
}
158155

159156
/* Error message */
160157
PyErr_Format(state->ZstdError,
161158
"Error when setting zstd %s parameter \"%s\", it "
162159
"should %d <= value <= %d, provided value is %d. "
163-
"(zstd v%s, %d-bit build)",
160+
"(%d-bit build)",
164161
type, name,
165162
bounds.lowerBound, bounds.upperBound, value_v,
166-
ZSTD_versionString(), 8*(int)sizeof(Py_ssize_t));
163+
8*(int)sizeof(Py_ssize_t));
167164
}
168165

169166
static inline _zstd_state*
@@ -558,150 +555,81 @@ static PyMethodDef _zstd_methods[] = {
558555
};
559556

560557

561-
#define ADD_INT_PREFIX_MACRO(module, macro) \
562-
do { \
563-
if (PyModule_AddIntConstant(module, "_" #macro, macro) < 0) { \
564-
return -1; \
565-
} \
566-
} while(0)
567-
568-
static int
569-
add_parameters(PyObject *module)
570-
{
571-
/* If add new parameters, please also add to cp_list/dp_list above. */
572-
573-
/* Compression parameters */
574-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_compressionLevel);
575-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_windowLog);
576-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_hashLog);
577-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_chainLog);
578-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_searchLog);
579-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_minMatch);
580-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_targetLength);
581-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_strategy);
582-
583-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_enableLongDistanceMatching);
584-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_ldmHashLog);
585-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_ldmMinMatch);
586-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_ldmBucketSizeLog);
587-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_ldmHashRateLog);
588-
589-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_contentSizeFlag);
590-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_checksumFlag);
591-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_dictIDFlag);
592-
593-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_nbWorkers);
594-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_jobSize);
595-
ADD_INT_PREFIX_MACRO(module, ZSTD_c_overlapLog);
596-
597-
/* Decompression parameters */
598-
ADD_INT_PREFIX_MACRO(module, ZSTD_d_windowLogMax);
599-
600-
/* ZSTD_strategy enum */
601-
ADD_INT_PREFIX_MACRO(module, ZSTD_fast);
602-
ADD_INT_PREFIX_MACRO(module, ZSTD_dfast);
603-
ADD_INT_PREFIX_MACRO(module, ZSTD_greedy);
604-
ADD_INT_PREFIX_MACRO(module, ZSTD_lazy);
605-
ADD_INT_PREFIX_MACRO(module, ZSTD_lazy2);
606-
ADD_INT_PREFIX_MACRO(module, ZSTD_btlazy2);
607-
ADD_INT_PREFIX_MACRO(module, ZSTD_btopt);
608-
ADD_INT_PREFIX_MACRO(module, ZSTD_btultra);
609-
ADD_INT_PREFIX_MACRO(module, ZSTD_btultra2);
610-
611-
return 0;
612-
}
613-
614-
static inline PyObject *
615-
get_zstd_version_info(void)
616-
{
617-
uint32_t ver = ZSTD_versionNumber();
618-
uint32_t major, minor, release;
619-
620-
major = ver / 10000;
621-
minor = (ver / 100) % 100;
622-
release = ver % 100;
623-
624-
return Py_BuildValue("III", major, minor, release);
625-
}
626-
627558
static inline int
628-
add_vars_to_module(PyObject *module)
559+
add_vars_to_module(PyObject *m)
629560
{
630-
PyObject *obj;
561+
#define ADD_INT_MACRO(MACRO) \
562+
if (PyModule_AddIntConstant((m), #MACRO, (MACRO)) < 0) { \
563+
return -1; \
564+
}
631565

632-
/* zstd_version, a str. */
633-
if (PyModule_AddStringConstant(module, "zstd_version",
634-
ZSTD_versionString()) < 0) {
566+
/* zstd_version_number, int */
567+
if (PyModule_AddIntConstant(m, "zstd_version_number",
568+
ZSTD_versionNumber()) < 0) {
635569
return -1;
636570
}
637571

638-
/* zstd_version_info, a tuple. */
639-
obj = get_zstd_version_info();
640-
if (PyModule_AddObjectRef(module, "zstd_version_info", obj) < 0) {
641-
Py_XDECREF(obj);
572+
/* zstd_version, str */
573+
if (PyModule_AddStringConstant(m, "zstd_version",
574+
ZSTD_versionString()) < 0) {
642575
return -1;
643576
}
644-
Py_DECREF(obj);
645577

646-
/* Add zstd parameters */
647-
if (add_parameters(module) < 0) {
578+
/* ZSTD_CLEVEL_DEFAULT, int */
579+
#if ZSTD_VERSION_NUMBER >= 10500
580+
if (PyModule_AddIntConstant(m, "ZSTD_CLEVEL_DEFAULT",
581+
ZSTD_defaultCLevel()) < 0) {
648582
return -1;
649583
}
650-
651-
/* _compressionLevel_values: (default, min, max)
652-
ZSTD_defaultCLevel() was added in zstd v1.5.0 */
653-
obj = Py_BuildValue("iii",
654-
#if ZSTD_VERSION_NUMBER < 10500
655-
ZSTD_CLEVEL_DEFAULT,
656584
#else
657-
ZSTD_defaultCLevel(),
585+
ADD_INT_MACRO(ZSTD_CLEVEL_DEFAULT);
658586
#endif
659-
ZSTD_minCLevel(),
660-
ZSTD_maxCLevel());
661-
if (PyModule_AddObjectRef(module,
662-
"_compressionLevel_values",
663-
obj) < 0) {
664-
Py_XDECREF(obj);
665-
return -1;
666-
}
667-
Py_DECREF(obj);
668587

669-
/* _ZSTD_CStreamSizes */
670-
obj = Py_BuildValue("II",
671-
(uint32_t)ZSTD_CStreamInSize(),
672-
(uint32_t)ZSTD_CStreamOutSize());
673-
if (PyModule_AddObjectRef(module, "_ZSTD_CStreamSizes", obj) < 0) {
674-
Py_XDECREF(obj);
588+
/* ZSTD_DStreamOutSize, int */
589+
if (PyModule_Add(m, "ZSTD_DStreamOutSize",
590+
PyLong_FromSize_t(ZSTD_DStreamOutSize())) < 0) {
675591
return -1;
676592
}
677-
Py_DECREF(obj);
678593

679-
/* _ZSTD_DStreamSizes */
680-
obj = Py_BuildValue("II",
681-
(uint32_t)ZSTD_DStreamInSize(),
682-
(uint32_t)ZSTD_DStreamOutSize());
683-
if (PyModule_AddObjectRef(module, "_ZSTD_DStreamSizes", obj) < 0) {
684-
Py_XDECREF(obj);
685-
return -1;
686-
}
687-
Py_DECREF(obj);
594+
/* Add zstd compression parameters. All should also be in cp_list. */
595+
ADD_INT_MACRO(ZSTD_c_compressionLevel);
596+
ADD_INT_MACRO(ZSTD_c_windowLog);
597+
ADD_INT_MACRO(ZSTD_c_hashLog);
598+
ADD_INT_MACRO(ZSTD_c_chainLog);
599+
ADD_INT_MACRO(ZSTD_c_searchLog);
600+
ADD_INT_MACRO(ZSTD_c_minMatch);
601+
ADD_INT_MACRO(ZSTD_c_targetLength);
602+
ADD_INT_MACRO(ZSTD_c_strategy);
688603

689-
/* _ZSTD_CONFIG */
690-
obj = Py_BuildValue("isOOO", 8*(int)sizeof(Py_ssize_t), "c",
691-
Py_False,
692-
Py_True,
693-
/* User mremap output buffer */
694-
#if defined(HAVE_MREMAP)
695-
Py_True
696-
#else
697-
Py_False
698-
#endif
699-
);
700-
if (PyModule_AddObjectRef(module, "_ZSTD_CONFIG", obj) < 0) {
701-
Py_XDECREF(obj);
702-
return -1;
703-
}
704-
Py_DECREF(obj);
604+
ADD_INT_MACRO(ZSTD_c_enableLongDistanceMatching);
605+
ADD_INT_MACRO(ZSTD_c_ldmHashLog);
606+
ADD_INT_MACRO(ZSTD_c_ldmMinMatch);
607+
ADD_INT_MACRO(ZSTD_c_ldmBucketSizeLog);
608+
ADD_INT_MACRO(ZSTD_c_ldmHashRateLog);
609+
610+
ADD_INT_MACRO(ZSTD_c_contentSizeFlag);
611+
ADD_INT_MACRO(ZSTD_c_checksumFlag);
612+
ADD_INT_MACRO(ZSTD_c_dictIDFlag);
613+
614+
ADD_INT_MACRO(ZSTD_c_nbWorkers);
615+
ADD_INT_MACRO(ZSTD_c_jobSize);
616+
ADD_INT_MACRO(ZSTD_c_overlapLog);
617+
618+
/* Add zstd decompression parameters. All should also be in dp_list. */
619+
ADD_INT_MACRO(ZSTD_d_windowLogMax);
620+
621+
/* ZSTD_strategy enum */
622+
ADD_INT_MACRO(ZSTD_fast);
623+
ADD_INT_MACRO(ZSTD_dfast);
624+
ADD_INT_MACRO(ZSTD_greedy);
625+
ADD_INT_MACRO(ZSTD_lazy);
626+
ADD_INT_MACRO(ZSTD_lazy2);
627+
ADD_INT_MACRO(ZSTD_btlazy2);
628+
ADD_INT_MACRO(ZSTD_btopt);
629+
ADD_INT_MACRO(ZSTD_btultra);
630+
ADD_INT_MACRO(ZSTD_btultra2);
631+
632+
#undef ADD_INT_MACRO
705633

706634
return 0;
707635
}

0 commit comments

Comments
 (0)