May 4, 2025 · Apr 26, 2025 · Apr 26, 2025 · Apr 26, 2025 · Apr 26, 2025 · Apr 26, 2025
diff --git a/Modules/_zstd/_zstdmodule.c b/Modules/_zstd/_zstdmodule.c
    Py_RETURN_NONE;
 }

 /*[clinic input]
 _zstd.compress

    data: Py_buffer
        A bytes-like object, data to be compressed.
    level: object  = None
        The compression level to use, defaults to ZSTD_CLEVEL_DEFAULT.
    options: object  = None
        A dict object that contains advanced compression parameters.
    zstd_dict: object  = None
        A ZstdDict object, a pre-trained zstd dictionary.

 Compress a block of data, return a bytes object of zstd compressed data.

 Refer to ZstdCompressor's docstring for a description of the
 optional arguments *level*, *options*, and *zstd_dict*.

 For incremental compression, use an ZstdCompressor instead.
 [clinic start generated code]*/

 static PyObject *
 _zstd_compress_impl(PyObject *module, Py_buffer *data, PyObject *level,
                    PyObject *options, PyObject *zstd_dict)
 /*[clinic end generated code: output=0cca9399ca5c95cc input=e8a7c59073af923c]*/
 {
    _zstd_state* const _module_state = PyModule_GetState(module);
    if (_module_state == NULL) {
        return NULL;
    }
    PyObject *ret = NULL;
    ZstdCompressor self = {0};

    /* Initialize & set ZstdCompressor */
    self.cctx = ZSTD_createCCtx();
    if (self.cctx == NULL) {
        PyErr_SetString(_module_state->ZstdError,
                        "Unable to create ZSTD_CCtx instance.");
        goto error;
    }

    if (level != Py_None && options != Py_None) {
        PyErr_SetString(PyExc_RuntimeError, "Only one of level or options should be used.");
        return NULL;
    }

    /* Set compressLevel/options to compression context */
    if (level != Py_None) {
        if (_PyZstd_set_c_parameters(&self, level, "level", "int") < 0) {
            goto error;
        }
    }

    if (options != Py_None) {
        if (_PyZstd_set_c_parameters(&self, options, "options", "dict") < 0) {
            goto error;
        }
    }

    /* Load dictionary to compression context */
    if (zstd_dict != Py_None) {
        if (_PyZstd_load_c_dict(&self, zstd_dict) < 0) {
            goto error;
        }
        self.dict = zstd_dict;
    }

    ret = compress_impl(&self, data, ZSTD_e_end);
    if (ret == NULL) {
        goto error;
    } else {
        goto success;
    }
 error:
    Py_CLEAR(ret);
 success:
    /* Free decompression context */
    ZSTD_freeCCtx(self.cctx);
    return ret;
 }

 /*[clinic input]
 _zstd.decompress

    data: Py_buffer
        A bytes-like object, zstd data to be decompressed.
    zstd_dict: object  = None
        A ZstdDict object, a pre-trained zstd dictionary.
    options: object  = None
        A dict object that contains advanced decompression parameters.

 Decompress one or more frames of data.

 Refer to ZstdDecompressor's docstring for a description of the
 optional arguments *zstd_dict*, *options*.

 For incremental decompression, use an ZstdDecompressor instead.
 [clinic start generated code]*/

 static PyObject *
 _zstd_decompress_impl(PyObject *module, Py_buffer *data, PyObject *zstd_dict,
                      PyObject *options)
 /*[clinic end generated code: output=2e8423588fb3b178 input=c18346e620e59039]*/
 {
    uint64_t decompressed_size;
    Py_ssize_t initial_size;
    ZstdDecompressor self = {0};
    ZSTD_inBuffer in;
    _zstd_state* const _module_state = PyModule_GetState(module);
    if (_module_state == NULL) {
        return NULL;
    }
    PyObject *ret = NULL;

    /* Initialize & set ZstdDecompressor */
    self.dctx = ZSTD_createDCtx();
    if (self.dctx == NULL) {
        PyErr_SetString(_module_state->ZstdError,
                        "Unable to create ZSTD_DCtx instance.");
        goto error;
    }
    self.at_frame_edge = 1;

    /* Load dictionary to decompression context */
    if (zstd_dict != Py_None) {
        if (_PyZstd_load_d_dict(&self, zstd_dict) < 0) {
            goto error;
        }
    }

    /* Set option to decompression context */
    if (options != Py_None) {
        if (_PyZstd_set_d_parameters(&self, options) < 0) {
            goto error;
        }
    }

    /* Prepare input data */
    in.src = data->buf;
    in.size = data->len;
    in.pos = 0;

    /* Get decompressed size */
    decompressed_size = ZSTD_getFrameContentSize(data->buf, data->len);
    /* These two zstd constants always > PY_SSIZE_T_MAX:
            ZSTD_CONTENTSIZE_UNKNOWN is (0ULL - 1)
            ZSTD_CONTENTSIZE_ERROR   is (0ULL - 2) */
    if (decompressed_size <= (uint64_t) PY_SSIZE_T_MAX) {
        initial_size = (Py_ssize_t) decompressed_size;
    } else {
        initial_size = -1;
    }

    /* Decompress */
    ret = decompress_impl(&self, &in, -1, initial_size,
                            TYPE_ENDLESS_DECOMPRESSOR);
    if (ret == NULL) {
        goto error;
    }

    /* Check data integrity. at_frame_edge flag is 1 when both the input and
        output streams are at a frame edge. */
    if (self.at_frame_edge == 0) {
        char *extra_msg = (Py_SIZE(ret) == 0) ? "." :
                            ", if want to output these decompressed data, use "
                            "the ZstdDecompressor class to decompress.";
        PyErr_Format(_module_state->ZstdError,
                        "Decompression failed: zstd data ends in an incomplete "
                        "frame, maybe the input data was truncated. Decompressed "
                        "data is %zd bytes%s",
                        Py_SIZE(ret), extra_msg);
        goto error;
    }

    goto success;

 error:
    Py_CLEAR(ret);
 success:
    /* Free decompression context */
    ZSTD_freeDCtx(self.dctx);
    return ret;
 }

 static PyMethodDef _zstd_methods[] = {
    _ZSTD__TRAIN_DICT_METHODDEF
    _ZSTD__FINALIZE_DICT_METHODDEF
    _ZSTD__GET_PARAM_BOUNDS_METHODDEF
    _ZSTD_GET_FRAME_SIZE_METHODDEF
    _ZSTD__GET_FRAME_INFO_METHODDEF
    _ZSTD__SET_PARAMETER_TYPES_METHODDEF
    _ZSTD_COMPRESS_METHODDEF
    _ZSTD_DECOMPRESS_METHODDEF

    {0}
 };

 static struct PyModuleDef_Slot _zstd_slots[] = {
    {Py_mod_exec, _zstd_exec},
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},

    {0}
 };
Original file line number	Diff line number	Diff line change
Expand Up		@@ -588,198 +588,13 @@ _zstd__set_parameter_types_impl(PyObject module, PyObject c_parameter_type,
		Py_RETURN_NONE;
		}

		/*[clinic input]
		_zstd.compress

		data: Py_buffer
		A bytes-like object, data to be compressed.
		level: object = None
		The compression level to use, defaults to ZSTD_CLEVEL_DEFAULT.
		options: object = None
		A dict object that contains advanced compression parameters.
		zstd_dict: object = None
		A ZstdDict object, a pre-trained zstd dictionary.

		Compress a block of data, return a bytes object of zstd compressed data.

		Refer to ZstdCompressor's docstring for a description of the
		optional arguments level, options, and zstd_dict.

		For incremental compression, use an ZstdCompressor instead.
		[clinic start generated code]*/

		static PyObject *
		_zstd_compress_impl(PyObject module, Py_buffer data, PyObject *level,
		PyObject options, PyObject zstd_dict)
		/[clinic end generated code: output=0cca9399ca5c95cc input=e8a7c59073af923c]/
		{
		_zstd_state* const _module_state = PyModule_GetState(module);
		if (_module_state == NULL) {
		return NULL;
		}
		PyObject *ret = NULL;
		ZstdCompressor self = {0};

		/* Initialize & set ZstdCompressor */
		self.cctx = ZSTD_createCCtx();
		if (self.cctx == NULL) {
		PyErr_SetString(_module_state->ZstdError,
		"Unable to create ZSTD_CCtx instance.");
		goto error;
		}

		if (level != Py_None && options != Py_None) {
		PyErr_SetString(PyExc_RuntimeError, "Only one of level or options should be used.");
		return NULL;
		}

		/* Set compressLevel/options to compression context */
		if (level != Py_None) {
		if (_PyZstd_set_c_parameters(&self, level, "level", "int") < 0) {
		goto error;
		}
		}

		if (options != Py_None) {
		if (_PyZstd_set_c_parameters(&self, options, "options", "dict") < 0) {
		goto error;
		}
		}

		/* Load dictionary to compression context */
		if (zstd_dict != Py_None) {
		if (_PyZstd_load_c_dict(&self, zstd_dict) < 0) {
		goto error;
		}
		self.dict = zstd_dict;
		}

		ret = compress_impl(&self, data, ZSTD_e_end);
		if (ret == NULL) {
		goto error;
		} else {
		goto success;
		}
		error:
		Py_CLEAR(ret);
		success:
		/* Free decompression context */
		ZSTD_freeCCtx(self.cctx);
		return ret;
		}

		/*[clinic input]
		_zstd.decompress

		data: Py_buffer
		A bytes-like object, zstd data to be decompressed.
		zstd_dict: object = None
		A ZstdDict object, a pre-trained zstd dictionary.
		options: object = None
		A dict object that contains advanced decompression parameters.

		Decompress one or more frames of data.

		Refer to ZstdDecompressor's docstring for a description of the
		optional arguments zstd_dict, options.

		For incremental decompression, use an ZstdDecompressor instead.
		[clinic start generated code]*/

		static PyObject *
		_zstd_decompress_impl(PyObject module, Py_buffer data, PyObject *zstd_dict,
		PyObject *options)
		/[clinic end generated code: output=2e8423588fb3b178 input=c18346e620e59039]/
		{
		uint64_t decompressed_size;
		Py_ssize_t initial_size;
		ZstdDecompressor self = {0};
		ZSTD_inBuffer in;
		_zstd_state* const _module_state = PyModule_GetState(module);
		if (_module_state == NULL) {
		return NULL;
		}
		PyObject *ret = NULL;

		/* Initialize & set ZstdDecompressor */
		self.dctx = ZSTD_createDCtx();
		if (self.dctx == NULL) {
		PyErr_SetString(_module_state->ZstdError,
		"Unable to create ZSTD_DCtx instance.");
		goto error;
		}
		self.at_frame_edge = 1;

		/* Load dictionary to decompression context */
		if (zstd_dict != Py_None) {
		if (_PyZstd_load_d_dict(&self, zstd_dict) < 0) {
		goto error;
		}
		}

		/* Set option to decompression context */
		if (options != Py_None) {
		if (_PyZstd_set_d_parameters(&self, options) < 0) {
		goto error;
		}
		}

		/* Prepare input data */
		in.src = data->buf;
		in.size = data->len;
		in.pos = 0;

		/* Get decompressed size */
		decompressed_size = ZSTD_getFrameContentSize(data->buf, data->len);
		/* These two zstd constants always > PY_SSIZE_T_MAX:
		ZSTD_CONTENTSIZE_UNKNOWN is (0ULL - 1)
		ZSTD_CONTENTSIZE_ERROR is (0ULL - 2) */
		if (decompressed_size <= (uint64_t) PY_SSIZE_T_MAX) {
		initial_size = (Py_ssize_t) decompressed_size;
		} else {
		initial_size = -1;
		}

		/* Decompress */
		ret = decompress_impl(&self, &in, -1, initial_size,
		TYPE_ENDLESS_DECOMPRESSOR);
		if (ret == NULL) {
		goto error;
		}

		/* Check data integrity. at_frame_edge flag is 1 when both the input and
		output streams are at a frame edge. */
		if (self.at_frame_edge == 0) {
		char *extra_msg = (Py_SIZE(ret) == 0) ? "." :
		", if want to output these decompressed data, use "
		"the ZstdDecompressor class to decompress.";
		PyErr_Format(_module_state->ZstdError,
		"Decompression failed: zstd data ends in an incomplete "
		"frame, maybe the input data was truncated. Decompressed "
		"data is %zd bytes%s",
		Py_SIZE(ret), extra_msg);
		goto error;
		}

		goto success;

		error:
		Py_CLEAR(ret);
		success:
		/* Free decompression context */
		ZSTD_freeDCtx(self.dctx);
		return ret;
		}

		static PyMethodDef _zstd_methods[] = {
		_ZSTD__TRAIN_DICT_METHODDEF
		_ZSTD__FINALIZE_DICT_METHODDEF
		_ZSTD__GET_PARAM_BOUNDS_METHODDEF
		_ZSTD_GET_FRAME_SIZE_METHODDEF
		_ZSTD__GET_FRAME_INFO_METHODDEF
		_ZSTD__SET_PARAMETER_TYPES_METHODDEF
		_ZSTD_COMPRESS_METHODDEF
		_ZSTD_DECOMPRESS_METHODDEF

		{0}
		};
Expand DownExpand Up		@@ -1130,6 +945,8 @@ _zstd_free(void *module)

		static struct PyModuleDef_Slot _zstd_slots[] = {
		{Py_mod_exec, _zstd_exec},
		{Py_mod_gil, Py_MOD_GIL_NOT_USED},

		{0}
		};

Expand Down