contrib/python-zstandard/c-ext/compressobj.c
changeset 37495 b1fb341d8a61
parent 31796 e0dc40530c5a
child 40121 73fef626dae3
--- a/contrib/python-zstandard/c-ext/compressobj.c	Sun Apr 08 01:08:43 2018 +0200
+++ b/contrib/python-zstandard/c-ext/compressobj.c	Mon Apr 09 10:13:29 2018 -0700
@@ -23,9 +23,13 @@
 	PyObject_Del(self);
 }
 
-static PyObject* ZstdCompressionObj_compress(ZstdCompressionObj* self, PyObject* args) {
-	const char* source;
-	Py_ssize_t sourceSize;
+static PyObject* ZstdCompressionObj_compress(ZstdCompressionObj* self, PyObject* args, PyObject* kwargs) {
+	static char* kwlist[] = {
+		"data",
+		NULL
+	};
+
+	Py_buffer source;
 	ZSTD_inBuffer input;
 	size_t zresult;
 	PyObject* result = NULL;
@@ -37,38 +41,43 @@
 	}
 
 #if PY_MAJOR_VERSION >= 3
-	if (!PyArg_ParseTuple(args, "y#:compress", &source, &sourceSize)) {
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:compress",
 #else
-	if (!PyArg_ParseTuple(args, "s#:compress", &source, &sourceSize)) {
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:compress",
 #endif
+		kwlist, &source)) {
 		return NULL;
 	}
 
-	input.src = source;
-	input.size = sourceSize;
+	if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {
+		PyErr_SetString(PyExc_ValueError,
+			"data buffer should be contiguous and have at most one dimension");
+		goto finally;
+	}
+
+	input.src = source.buf;
+	input.size = source.len;
 	input.pos = 0;
 
-	while ((ssize_t)input.pos < sourceSize) {
+	while ((ssize_t)input.pos < source.len) {
 		Py_BEGIN_ALLOW_THREADS
-		if (self->compressor->mtcctx) {
-			zresult = ZSTDMT_compressStream(self->compressor->mtcctx,
-				&self->output, &input);
-		}
-		else {
-			zresult = ZSTD_compressStream(self->compressor->cstream, &self->output, &input);
-		}
+			zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
+				&input, ZSTD_e_continue);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
 			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
-			return NULL;
+			Py_CLEAR(result);
+			goto finally;
 		}
 
 		if (self->output.pos) {
 			if (result) {
 				resultSize = PyBytes_GET_SIZE(result);
-				if (-1 == _PyBytes_Resize(&result, resultSize + self->output.pos)) {
-					return NULL;
+
+				if (safe_pybytes_resize(&result, resultSize + self->output.pos)) {
+					Py_CLEAR(result);
+					goto finally;
 				}
 
 				memcpy(PyBytes_AS_STRING(result) + resultSize,
@@ -77,7 +86,7 @@
 			else {
 				result = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
 				if (!result) {
-					return NULL;
+					goto finally;
 				}
 			}
 
@@ -85,21 +94,29 @@
 		}
 	}
 
-	if (result) {
-		return result;
+	if (NULL == result) {
+		result = PyBytes_FromString("");
 	}
-	else {
-		return PyBytes_FromString("");
-	}
+
+finally:
+	PyBuffer_Release(&source);
+
+	return result;
 }
 
-static PyObject* ZstdCompressionObj_flush(ZstdCompressionObj* self, PyObject* args) {
+static PyObject* ZstdCompressionObj_flush(ZstdCompressionObj* self, PyObject* args, PyObject* kwargs) {
+	static char* kwlist[] = {
+		"flush_mode",
+		NULL
+	};
+
 	int flushMode = compressorobj_flush_finish;
 	size_t zresult;
 	PyObject* result = NULL;
 	Py_ssize_t resultSize = 0;
+	ZSTD_inBuffer input;
 
-	if (!PyArg_ParseTuple(args, "|i:flush", &flushMode)) {
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:flush", kwlist, &flushMode)) {
 		return NULL;
 	}
 
@@ -115,16 +132,16 @@
 
 	assert(self->output.pos == 0);
 
+	input.src = NULL;
+	input.size = 0;
+	input.pos = 0;
+
 	if (flushMode == compressorobj_flush_block) {
 		/* The output buffer is of size ZSTD_CStreamOutSize(), which is 
 		   guaranteed to hold a full block. */
 		Py_BEGIN_ALLOW_THREADS
-		if (self->compressor->mtcctx) {
-			zresult = ZSTDMT_flushStream(self->compressor->mtcctx, &self->output);
-		}
-		else {
-			zresult = ZSTD_flushStream(self->compressor->cstream, &self->output);
-		}
+			zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
+				&input, ZSTD_e_flush);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
@@ -156,12 +173,8 @@
 	self->finished = 1;
 
 	while (1) {
-		if (self->compressor->mtcctx) {
-			zresult = ZSTDMT_endStream(self->compressor->mtcctx, &self->output);
-		}
-		else {
-			zresult = ZSTD_endStream(self->compressor->cstream, &self->output);
-		}
+		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
+			&input, ZSTD_e_end);
 		if (ZSTD_isError(zresult)) {
 			PyErr_Format(ZstdError, "error ending compression stream: %s",
 				ZSTD_getErrorName(zresult));
@@ -171,7 +184,9 @@
 		if (self->output.pos) {
 			if (result) {
 				resultSize = PyBytes_GET_SIZE(result);
-				if (-1 == _PyBytes_Resize(&result, resultSize + self->output.pos)) {
+
+				if (safe_pybytes_resize(&result, resultSize + self->output.pos)) {
+					Py_XDECREF(result);
 					return NULL;
 				}
 
@@ -202,9 +217,9 @@
 }
 
 static PyMethodDef ZstdCompressionObj_methods[] = {
-	{ "compress", (PyCFunction)ZstdCompressionObj_compress, METH_VARARGS,
+	{ "compress", (PyCFunction)ZstdCompressionObj_compress, METH_VARARGS | METH_KEYWORDS,
 	PyDoc_STR("compress data") },
-	{ "flush", (PyCFunction)ZstdCompressionObj_flush, METH_VARARGS,
+	{ "flush", (PyCFunction)ZstdCompressionObj_flush, METH_VARARGS | METH_KEYWORDS,
 	PyDoc_STR("finish compression operation") },
 	{ NULL, NULL }
 };