comparison contrib/python-zstandard/c-ext/compressor.c @ 30822:b54a2984cdd4

zstd: vendor python-zstandard 0.6.0

Commit 63c68d6f5fc8de4afd9bde81b13b537beb4e47e8 from
https://github.com/indygreg/python-zstandard is imported without
modifications (other than removing unwanted files).

This includes minor performance and feature improvements. It also
changes the vendored zstd library from 1.1.1 to 1.1.2.

# no-check-commit

30821:7005c03f7387

*/

#include "python-zstandard.h"

extern PyObject* ZstdError;

/**
* Initialize a zstd CStream from a ZstdCompressor instance.
*
* Returns a ZSTD_CStream on success or NULL on failure. If NULL, a Python

		return NULL;
	}

	return cstream;
}


PyDoc_STRVAR(ZstdCompressor__doc__,
"ZstdCompressor(level=None, dict_data=None, compression_params=None)\n"
"\n"
"Create an object used to perform Zstandard compression.\n"

	CompressionParametersObject* params = NULL;
	PyObject* writeChecksum = NULL;
	PyObject* writeContentSize = NULL;
	PyObject* writeDictID = NULL;

	self->dict = NULL;
	self->cparams = NULL;
	self->cdict = NULL;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!O!OOO", kwlist,

		PyErr_Format(PyExc_ValueError, "level must be less than %d",
			ZSTD_maxCLevel() + 1);
		return -1;
	}

	self->compressionLevel = level;

	if (dict) {
		self->dict = dict;
		Py_INCREF(dict);

	Py_XDECREF(self->dict);

	if (self->cdict) {
		ZSTD_freeCDict(self->cdict);
		self->cdict = NULL;
	}

	PyObject_Del(self);
}

	return res;
}

PyDoc_STRVAR(ZstdCompressor_compress__doc__,
"compress(data)\n"
"\n"
"Compress data in a single operation.\n"
"\n"
"This is the simplest mechanism to perform compression: simply pass in a\n"
"value and get a compressed value back. It is almost the most prone to abuse.\n"
"The input and output values must fit in memory, so passing in very large\n"
"values can result in excessive memory usage. For this reason, one of the\n"
"streaming based APIs is preferred for larger values.\n"
);

static PyObject* ZstdCompressor_compress(ZstdCompressor* self, PyObject* args) {
	const char* source;
	Py_ssize_t sourceSize;
	size_t destSize;
	ZSTD_CCtx* cctx;
	PyObject* output;
	char* dest;
	void* dictData = NULL;
	size_t dictSize = 0;
	size_t zresult;
	ZSTD_parameters zparams;
	ZSTD_customMem zmem;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTuple(args, "y#", &source, &sourceSize)) {
#else
	if (!PyArg_ParseTuple(args, "s#", &source, &sourceSize)) {
#endif
		return NULL;
	}

	destSize = ZSTD_compressBound(sourceSize);
	output = PyBytes_FromStringAndSize(NULL, destSize);
	if (!output) {
		return NULL;
	}

	dest = PyBytes_AsString(output);

	cctx = ZSTD_createCCtx();
	if (!cctx) {
		Py_DECREF(output);
		PyErr_SetString(ZstdError, "could not create CCtx");
		return NULL;
	}

	if (self->dict) {
		dictData = self->dict->dictData;
		dictSize = self->dict->dictSize;
	}

	zparams.fParams = self->fparams;

	/* The raw dict data has to be processed before it can be used. Since this
	adds overhead - especially if multiple dictionary compression operations
	are performed on the same ZstdCompressor instance - we create a
	ZSTD_CDict once and reuse it for all operations. */

	/* TODO the zparams (which can be derived from the source data size) used
	on first invocation are effectively reused for subsequent operations. This
	may not be appropriate if input sizes vary significantly and could affect
	chosen compression parameters.
	https://github.com/facebook/zstd/issues/358 tracks this issue. */
	if (dictData && !self->cdict) {
		Py_BEGIN_ALLOW_THREADS
		memset(&zmem, 0, sizeof(zmem));
		self->cdict = ZSTD_createCDict_advanced(dictData, dictSize, zparams, zmem);
		Py_END_ALLOW_THREADS

		if (!self->cdict) {
			Py_DECREF(output);
			ZSTD_freeCCtx(cctx);
			PyErr_SetString(ZstdError, "could not create compression dictionary");
			return NULL;
		}
	}

	Py_BEGIN_ALLOW_THREADS
	/* By avoiding ZSTD_compress(), we don't necessarily write out content
	   size. This means the argument to ZstdCompressor to control frame
	   parameters is honored. */
	if (self->cdict) {
		zresult = ZSTD_compress_usingCDict(cctx, dest, destSize,
			source, sourceSize, self->cdict);
	}
	else {
		zresult = ZSTD_compress_advanced(cctx, dest, destSize,
			source, sourceSize, dictData, dictSize, zparams);
	}
	Py_END_ALLOW_THREADS

	ZSTD_freeCCtx(cctx);

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "cannot compress: %s", ZSTD_getErrorName(zresult));
		Py_CLEAR(output);
		return NULL;

	result->output.pos = 0;

	result->compressor = self;
	Py_INCREF(result->compressor);

	result->flushed = 0;

	return result;
}

PyDoc_STRVAR(ZstdCompressor_read_from__doc__,

	return result;
}

static PyMethodDef ZstdCompressor_methods[] = {
	{ "compress", (PyCFunction)ZstdCompressor_compress, METH_VARARGS,
	ZstdCompressor_compress__doc__ },
	{ "compressobj", (PyCFunction)ZstdCompressor_compressobj,
	METH_VARARGS | METH_KEYWORDS, ZstdCompressionObj__doc__ },
	{ "copy_stream", (PyCFunction)ZstdCompressor_copy_stream,
	METH_VARARGS | METH_KEYWORDS, ZstdCompressor_copy_stream__doc__ },
	{ "read_from", (PyCFunction)ZstdCompressor_read_from,

30822:b54a2984cdd4

*/

#include "python-zstandard.h"

extern PyObject* ZstdError;

int populate_cdict(ZstdCompressor* compressor, void* dictData, size_t dictSize, ZSTD_parameters* zparams) {
	ZSTD_customMem zmem;
	assert(!compressor->cdict);
	Py_BEGIN_ALLOW_THREADS
	memset(&zmem, 0, sizeof(zmem));
	compressor->cdict = ZSTD_createCDict_advanced(compressor->dict->dictData,
		compressor->dict->dictSize, *zparams, zmem);
	Py_END_ALLOW_THREADS

	if (!compressor->cdict) {
		PyErr_SetString(ZstdError, "could not create compression dictionary");
		return 1;
	}

	return 0;
}

/**
* Initialize a zstd CStream from a ZstdCompressor instance.
*
* Returns a ZSTD_CStream on success or NULL on failure. If NULL, a Python

		return NULL;
	}

	return cstream;
}

PyDoc_STRVAR(ZstdCompressor__doc__,
"ZstdCompressor(level=None, dict_data=None, compression_params=None)\n"
"\n"
"Create an object used to perform Zstandard compression.\n"

	CompressionParametersObject* params = NULL;
	PyObject* writeChecksum = NULL;
	PyObject* writeContentSize = NULL;
	PyObject* writeDictID = NULL;

	self->cctx = NULL;
	self->dict = NULL;
	self->cparams = NULL;
	self->cdict = NULL;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!O!OOO", kwlist,

		PyErr_Format(PyExc_ValueError, "level must be less than %d",
			ZSTD_maxCLevel() + 1);
		return -1;
	}

	/* We create a ZSTD_CCtx for reuse among multiple operations to reduce the
	   overhead of each compression operation. */
	self->cctx = ZSTD_createCCtx();
	if (!self->cctx) {
		PyErr_NoMemory();
		return -1;
	}

	self->compressionLevel = level;

	if (dict) {
		self->dict = dict;
		Py_INCREF(dict);

	Py_XDECREF(self->dict);

	if (self->cdict) {
		ZSTD_freeCDict(self->cdict);
		self->cdict = NULL;
	}

	if (self->cctx) {
		ZSTD_freeCCtx(self->cctx);
		self->cctx = NULL;
	}

	PyObject_Del(self);
}

	return res;
}

PyDoc_STRVAR(ZstdCompressor_compress__doc__,
"compress(data, allow_empty=False)\n"
"\n"
"Compress data in a single operation.\n"
"\n"
"This is the simplest mechanism to perform compression: simply pass in a\n"
"value and get a compressed value back. It is almost the most prone to abuse.\n"
"The input and output values must fit in memory, so passing in very large\n"
"values can result in excessive memory usage. For this reason, one of the\n"
"streaming based APIs is preferred for larger values.\n"
);

static PyObject* ZstdCompressor_compress(ZstdCompressor* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"data",
		"allow_empty",
		NULL
	};

	const char* source;
	Py_ssize_t sourceSize;
	PyObject* allowEmpty = NULL;
	size_t destSize;
	PyObject* output;
	char* dest;
	void* dictData = NULL;
	size_t dictSize = 0;
	size_t zresult;
	ZSTD_parameters zparams;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y#|O",
#else
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s#|O",
#endif
		kwlist, &source, &sourceSize, &allowEmpty)) {
		return NULL;
	}

	/* Limitation in zstd C API doesn't let decompression side distinguish
	   between content size of 0 and unknown content size. This can make round
	   tripping via Python difficult. Until this is fixed, require a flag
	   to fire the footgun.
	   https://github.com/indygreg/python-zstandard/issues/11 */
	if (0 == sourceSize && self->fparams.contentSizeFlag
		&& (!allowEmpty || PyObject_Not(allowEmpty))) {
		PyErr_SetString(PyExc_ValueError, "cannot write empty inputs when writing content sizes");
		return NULL;
	}

	destSize = ZSTD_compressBound(sourceSize);
	output = PyBytes_FromStringAndSize(NULL, destSize);
	if (!output) {
		return NULL;
	}

	dest = PyBytes_AsString(output);

	if (self->dict) {
		dictData = self->dict->dictData;
		dictSize = self->dict->dictSize;
	}

	zparams.fParams = self->fparams;

	/* The raw dict data has to be processed before it can be used. Since this
	adds overhead - especially if multiple dictionary compression operations
	are performed on the same ZstdCompressor instance - we create a
	ZSTD_CDict once and reuse it for all operations.

	Note: the compression parameters used for the first invocation (possibly
	derived from the source size) will be reused on all subsequent invocations.
	https://github.com/facebook/zstd/issues/358 contains more info. We could
	potentially add an argument somewhere to control this behavior.
	*/
	if (dictData && !self->cdict) {
		if (populate_cdict(self, dictData, dictSize, &zparams)) {
			Py_DECREF(output);
			return NULL;
		}
	}

	Py_BEGIN_ALLOW_THREADS
	/* By avoiding ZSTD_compress(), we don't necessarily write out content
	   size. This means the argument to ZstdCompressor to control frame
	   parameters is honored. */
	if (self->cdict) {
		zresult = ZSTD_compress_usingCDict(self->cctx, dest, destSize,
			source, sourceSize, self->cdict);
	}
	else {
		zresult = ZSTD_compress_advanced(self->cctx, dest, destSize,
			source, sourceSize, dictData, dictSize, zparams);
	}
	Py_END_ALLOW_THREADS

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "cannot compress: %s", ZSTD_getErrorName(zresult));
		Py_CLEAR(output);
		return NULL;

	result->output.pos = 0;

	result->compressor = self;
	Py_INCREF(result->compressor);

	result->finished = 0;

	return result;
}

PyDoc_STRVAR(ZstdCompressor_read_from__doc__,

	return result;
}

static PyMethodDef ZstdCompressor_methods[] = {
	{ "compress", (PyCFunction)ZstdCompressor_compress,
	METH_VARARGS | METH_KEYWORDS, ZstdCompressor_compress__doc__ },
	{ "compressobj", (PyCFunction)ZstdCompressor_compressobj,
	METH_VARARGS | METH_KEYWORDS, ZstdCompressionObj__doc__ },
	{ "copy_stream", (PyCFunction)ZstdCompressor_copy_stream,
	METH_VARARGS | METH_KEYWORDS, ZstdCompressor_copy_stream__doc__ },
	{ "read_from", (PyCFunction)ZstdCompressor_read_from,