aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/Modules/_zstd/compressor.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/_zstd/compressor.c')
-rw-r--r--Modules/_zstd/compressor.c222
1 files changed, 121 insertions, 101 deletions
diff --git a/Modules/_zstd/compressor.c b/Modules/_zstd/compressor.c
index 8f934858ef7..0fc3d7d36c6 100644
--- a/Modules/_zstd/compressor.c
+++ b/Modules/_zstd/compressor.c
@@ -49,98 +49,106 @@ typedef struct {
#include "clinic/compressor.c.h"
static int
-_zstd_set_c_parameters(ZstdCompressor *self, PyObject *level_or_options,
- const char *arg_name, const char* arg_type)
+_zstd_set_c_level(ZstdCompressor *self, int level)
{
- size_t zstd_ret;
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
- if (mod_state == NULL) {
+ /* Set integer compression level */
+ int min_level = ZSTD_minCLevel();
+ int max_level = ZSTD_maxCLevel();
+ if (level < min_level || level > max_level) {
+ PyErr_Format(PyExc_ValueError,
+ "illegal compression level %d; the valid range is [%d, %d]",
+ level, min_level, max_level);
return -1;
}
- /* Integer compression level */
- if (PyLong_Check(level_or_options)) {
- int level = PyLong_AsInt(level_or_options);
- if (level == -1 && PyErr_Occurred()) {
- PyErr_Format(PyExc_ValueError,
- "Compression level should be an int value between %d and %d.",
- ZSTD_minCLevel(), ZSTD_maxCLevel());
+ /* Save for generating ZSTD_CDICT */
+ self->compression_level = level;
+
+ /* Set compressionLevel to compression context */
+ size_t zstd_ret = ZSTD_CCtx_setParameter(
+ self->cctx, ZSTD_c_compressionLevel, level);
+
+ /* Check error */
+ if (ZSTD_isError(zstd_ret)) {
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
+ if (mod_state == NULL) {
return -1;
}
+ set_zstd_error(mod_state, ERR_SET_C_LEVEL, zstd_ret);
+ return -1;
+ }
+ return 0;
+}
- /* Save for generating ZSTD_CDICT */
- self->compression_level = level;
+static int
+_zstd_set_c_parameters(ZstdCompressor *self, PyObject *options)
+{
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
+ if (mod_state == NULL) {
+ return -1;
+ }
- /* Set compressionLevel to compression context */
- zstd_ret = ZSTD_CCtx_setParameter(self->cctx,
- ZSTD_c_compressionLevel,
- level);
+ if (!PyDict_Check(options)) {
+ PyErr_Format(PyExc_TypeError,
+ "ZstdCompressor() argument 'options' must be dict, not %T",
+ options);
+ return -1;
+ }
- /* Check error */
- if (ZSTD_isError(zstd_ret)) {
- set_zstd_error(mod_state, ERR_SET_C_LEVEL, zstd_ret);
+ Py_ssize_t pos = 0;
+ PyObject *key, *value;
+ while (PyDict_Next(options, &pos, &key, &value)) {
+ /* Check key type */
+ if (Py_TYPE(key) == mod_state->DParameter_type) {
+ PyErr_SetString(PyExc_TypeError,
+ "compression options dictionary key must not be a "
+ "DecompressionParameter attribute");
return -1;
}
- return 0;
- }
- /* Options dict */
- if (PyDict_Check(level_or_options)) {
- PyObject *key, *value;
- Py_ssize_t pos = 0;
+ Py_INCREF(key);
+ Py_INCREF(value);
+ int key_v = PyLong_AsInt(key);
+ Py_DECREF(key);
+ if (key_v == -1 && PyErr_Occurred()) {
+ Py_DECREF(value);
+ return -1;
+ }
- while (PyDict_Next(level_or_options, &pos, &key, &value)) {
- /* Check key type */
- if (Py_TYPE(key) == mod_state->DParameter_type) {
- PyErr_SetString(PyExc_TypeError,
- "Key of compression option dict should "
- "NOT be DecompressionParameter.");
- return -1;
- }
+ int value_v = PyLong_AsInt(value);
+ Py_DECREF(value);
+ if (value_v == -1 && PyErr_Occurred()) {
+ return -1;
+ }
- int key_v = PyLong_AsInt(key);
- if (key_v == -1 && PyErr_Occurred()) {
- PyErr_SetString(PyExc_ValueError,
- "Key of options dict should be a CompressionParameter attribute.");
+ if (key_v == ZSTD_c_compressionLevel) {
+ if (_zstd_set_c_level(self, value_v) < 0) {
return -1;
}
-
- // TODO(emmatyping): check bounds when there is a value error here for better
- // error message?
- int value_v = PyLong_AsInt(value);
- if (value_v == -1 && PyErr_Occurred()) {
- PyErr_SetString(PyExc_ValueError,
- "Value of option dict should be an int.");
- return -1;
+ continue;
+ }
+ if (key_v == ZSTD_c_nbWorkers) {
+ /* From the zstd library docs:
+ 1. When nbWorkers >= 1, triggers asynchronous mode when
+ used with ZSTD_compressStream2().
+ 2, Default value is `0`, aka "single-threaded mode" : no
+ worker is spawned, compression is performed inside
+ caller's thread, all invocations are blocking. */
+ if (value_v != 0) {
+ self->use_multithread = 1;
}
+ }
- if (key_v == ZSTD_c_compressionLevel) {
- /* Save for generating ZSTD_CDICT */
- self->compression_level = value_v;
- }
- else if (key_v == ZSTD_c_nbWorkers) {
- /* From the zstd library docs:
- 1. When nbWorkers >= 1, triggers asynchronous mode when
- used with ZSTD_compressStream2().
- 2, Default value is `0`, aka "single-threaded mode" : no
- worker is spawned, compression is performed inside
- caller's thread, all invocations are blocking. */
- if (value_v != 0) {
- self->use_multithread = 1;
- }
- }
+ /* Set parameter to compression context */
+ size_t zstd_ret = ZSTD_CCtx_setParameter(self->cctx, key_v, value_v);
- /* Set parameter to compression context */
- zstd_ret = ZSTD_CCtx_setParameter(self->cctx, key_v, value_v);
- if (ZSTD_isError(zstd_ret)) {
- set_parameter_error(mod_state, 1, key_v, value_v);
- return -1;
- }
+ /* Check error */
+ if (ZSTD_isError(zstd_ret)) {
+ set_parameter_error(1, key_v, value_v);
+ return -1;
}
- return 0;
}
- PyErr_Format(PyExc_TypeError, "Invalid type for %s. Expected %s", arg_name, arg_type);
- return -1;
+ return 0;
}
static void
@@ -173,16 +181,13 @@ _get_CDict(ZstdDict *self, int compressionLevel)
}
if (capsule == NULL) {
/* Create ZSTD_CDict instance */
- char *dict_buffer = PyBytes_AS_STRING(self->dict_content);
- Py_ssize_t dict_len = Py_SIZE(self->dict_content);
Py_BEGIN_ALLOW_THREADS
- cdict = ZSTD_createCDict(dict_buffer,
- dict_len,
+ cdict = ZSTD_createCDict(self->dict_buffer, self->dict_len,
compressionLevel);
Py_END_ALLOW_THREADS
if (cdict == NULL) {
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
if (mod_state != NULL) {
PyErr_SetString(mod_state->ZstdError,
"Failed to create a ZSTD_CDict instance from "
@@ -236,17 +241,13 @@ _zstd_load_impl(ZstdCompressor *self, ZstdDict *zd,
else if (type == DICT_TYPE_UNDIGESTED) {
/* Load a dictionary.
It doesn't override compression context's parameters. */
- zstd_ret = ZSTD_CCtx_loadDictionary(
- self->cctx,
- PyBytes_AS_STRING(zd->dict_content),
- Py_SIZE(zd->dict_content));
+ zstd_ret = ZSTD_CCtx_loadDictionary(self->cctx, zd->dict_buffer,
+ zd->dict_len);
}
else if (type == DICT_TYPE_PREFIX) {
/* Load a prefix */
- zstd_ret = ZSTD_CCtx_refPrefix(
- self->cctx,
- PyBytes_AS_STRING(zd->dict_content),
- Py_SIZE(zd->dict_content));
+ zstd_ret = ZSTD_CCtx_refPrefix(self->cctx, zd->dict_buffer,
+ zd->dict_len);
}
else {
Py_UNREACHABLE();
@@ -263,7 +264,7 @@ _zstd_load_impl(ZstdCompressor *self, ZstdDict *zd,
static int
_zstd_load_c_dict(ZstdCompressor *self, PyObject *dict)
{
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
if (mod_state == NULL) {
return -1;
}
@@ -296,9 +297,9 @@ _zstd_load_c_dict(ZstdCompressor *self, PyObject *dict)
else if (ret > 0) {
/* type == -1 may indicate an error. */
type = PyLong_AsInt(PyTuple_GET_ITEM(dict, 1));
- if (type == DICT_TYPE_DIGESTED ||
- type == DICT_TYPE_UNDIGESTED ||
- type == DICT_TYPE_PREFIX)
+ if (type == DICT_TYPE_DIGESTED
+ || type == DICT_TYPE_UNDIGESTED
+ || type == DICT_TYPE_PREFIX)
{
assert(type >= 0);
zd = (ZstdDict*)PyTuple_GET_ITEM(dict, 0);
@@ -349,7 +350,7 @@ _zstd_ZstdCompressor_new_impl(PyTypeObject *type, PyObject *level,
/* Compression context */
self->cctx = ZSTD_createCCtx();
if (self->cctx == NULL) {
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
if (mod_state != NULL) {
PyErr_SetString(mod_state->ZstdError,
"Unable to create ZSTD_CCtx instance.");
@@ -361,19 +362,35 @@ _zstd_ZstdCompressor_new_impl(PyTypeObject *type, PyObject *level,
self->last_mode = ZSTD_e_end;
if (level != Py_None && options != Py_None) {
- PyErr_SetString(PyExc_RuntimeError, "Only one of level or options should be used.");
+ PyErr_SetString(PyExc_TypeError,
+ "Only one of level or options should be used.");
goto error;
}
- /* Set compressLevel/options to compression context */
+ /* Set compression level */
if (level != Py_None) {
- if (_zstd_set_c_parameters(self, level, "level", "int") < 0) {
+ if (!PyLong_Check(level)) {
+ PyErr_SetString(PyExc_TypeError,
+ "invalid type for level, expected int");
+ goto error;
+ }
+ int level_v = PyLong_AsInt(level);
+ if (level_v == -1 && PyErr_Occurred()) {
+ if (PyErr_ExceptionMatches(PyExc_OverflowError)) {
+ PyErr_Format(PyExc_ValueError,
+ "illegal compression level; the valid range is [%d, %d]",
+ ZSTD_minCLevel(), ZSTD_maxCLevel());
+ }
+ goto error;
+ }
+ if (_zstd_set_c_level(self, level_v) < 0) {
goto error;
}
}
+ /* Set options dictionary */
if (options != Py_None) {
- if (_zstd_set_c_parameters(self, options, "options", "dict") < 0) {
+ if (_zstd_set_c_parameters(self, options) < 0) {
goto error;
}
}
@@ -450,7 +467,7 @@ compress_lock_held(ZstdCompressor *self, Py_buffer *data,
}
if (_OutputBuffer_InitWithSize(&buffer, &out, -1,
- (Py_ssize_t) output_buffer_size) < 0) {
+ (Py_ssize_t) output_buffer_size) < 0) {
goto error;
}
@@ -463,7 +480,7 @@ compress_lock_held(ZstdCompressor *self, Py_buffer *data,
/* Check error */
if (ZSTD_isError(zstd_ret)) {
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
if (mod_state != NULL) {
set_zstd_error(mod_state, ERR_COMPRESS, zstd_ret);
}
@@ -526,13 +543,16 @@ compress_mt_continue_lock_held(ZstdCompressor *self, Py_buffer *data)
while (1) {
Py_BEGIN_ALLOW_THREADS
do {
- zstd_ret = ZSTD_compressStream2(self->cctx, &out, &in, ZSTD_e_continue);
- } while (out.pos != out.size && in.pos != in.size && !ZSTD_isError(zstd_ret));
+ zstd_ret = ZSTD_compressStream2(self->cctx, &out, &in,
+ ZSTD_e_continue);
+ } while (out.pos != out.size
+ && in.pos != in.size
+ && !ZSTD_isError(zstd_ret));
Py_END_ALLOW_THREADS
/* Check error */
if (ZSTD_isError(zstd_ret)) {
- _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
+ _zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));
if (mod_state != NULL) {
set_zstd_error(mod_state, ERR_COMPRESS, zstd_ret);
}
@@ -679,12 +699,12 @@ static PyMethodDef ZstdCompressor_methods[] = {
PyDoc_STRVAR(ZstdCompressor_last_mode_doc,
"The last mode used to this compressor object, its value can be .CONTINUE,\n"
".FLUSH_BLOCK, .FLUSH_FRAME. Initialized to .FLUSH_FRAME.\n\n"
-"It can be used to get the current state of a compressor, such as, data flushed,\n"
-"a frame ended.");
+"It can be used to get the current state of a compressor, such as, data\n"
+"flushed, or a frame ended.");
static PyMemberDef ZstdCompressor_members[] = {
{"last_mode", Py_T_INT, offsetof(ZstdCompressor, last_mode),
- Py_READONLY, ZstdCompressor_last_mode_doc},
+ Py_READONLY, ZstdCompressor_last_mode_doc},
{NULL}
};