contrib/python-zstandard/c-ext/decompressionwriter.c

--- a/contrib/python-zstandard/c-ext/decompressionwriter.c	Sun Apr 08 01:08:43 2018 +0200
+++ b/contrib/python-zstandard/c-ext/decompressionwriter.c	Mon Apr 09 10:13:29 2018 -0700
@@ -27,7 +27,7 @@
 		return NULL;
 	}
 
-	if (0 != init_dstream(self->decompressor)) {
+	if (ensure_dctx(self->decompressor, 1)) {
 		return NULL;
 	}
 
@@ -44,18 +44,17 @@
 }
 
 static PyObject* ZstdDecompressionWriter_memory_size(ZstdDecompressionWriter* self) {
-	if (!self->decompressor->dstream) {
-		PyErr_SetString(ZstdError, "cannot determine size of inactive decompressor; "
-			"call when context manager is active");
-		return NULL;
-	}
-
-	return PyLong_FromSize_t(ZSTD_sizeof_DStream(self->decompressor->dstream));
+	return PyLong_FromSize_t(ZSTD_sizeof_DCtx(self->decompressor->dctx));
 }
 
-static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, PyObject* args) {
-	const char* source;
-	Py_ssize_t sourceSize;
+static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, PyObject* args, PyObject* kwargs) {
+	static char* kwlist[] = {
+		"data",
+		NULL
+	};
+
+	PyObject* result = NULL;
+	Py_buffer source;
 	size_t zresult = 0;
 	ZSTD_inBuffer input;
 	ZSTD_outBuffer output;
@@ -63,41 +62,47 @@
 	Py_ssize_t totalWrite = 0;
 
 #if PY_MAJOR_VERSION >= 3
-	if (!PyArg_ParseTuple(args, "y#:write", &source, &sourceSize)) {
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:write",
 #else
-	if (!PyArg_ParseTuple(args, "s#:write", &source, &sourceSize)) {
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:write",
 #endif
+		kwlist, &source)) {
 		return NULL;
 	}
 
+	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;
+	}
+
 	if (!self->entered) {
 		PyErr_SetString(ZstdError, "write must be called from an active context manager");
-		return NULL;
+		goto finally;
 	}
 
-	assert(self->decompressor->dstream);
-
 	output.dst = PyMem_Malloc(self->outSize);
 	if (!output.dst) {
-		return PyErr_NoMemory();
+		PyErr_NoMemory();
+		goto finally;
 	}
 	output.size = self->outSize;
 	output.pos = 0;
 
-	input.src = source;
-	input.size = sourceSize;
+	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
-		zresult = ZSTD_decompressStream(self->decompressor->dstream, &output, &input);
+		zresult = ZSTD_decompress_generic(self->decompressor->dctx, &output, &input);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
 			PyMem_Free(output.dst);
 			PyErr_Format(ZstdError, "zstd decompress error: %s",
 				ZSTD_getErrorName(zresult));
-			return NULL;
+			goto finally;
 		}
 
 		if (output.pos) {
@@ -115,7 +120,11 @@
 
 	PyMem_Free(output.dst);
 
-	return PyLong_FromSsize_t(totalWrite);
+	result = PyLong_FromSsize_t(totalWrite);
+
+finally:
+	PyBuffer_Release(&source);
+	return result;
 }
 
 static PyMethodDef ZstdDecompressionWriter_methods[] = {
@@ -125,7 +134,7 @@
 	PyDoc_STR("Exit a decompression context.") },
 	{ "memory_size", (PyCFunction)ZstdDecompressionWriter_memory_size, METH_NOARGS,
 	PyDoc_STR("Obtain the memory size in bytes of the underlying decompressor.") },
-	{ "write", (PyCFunction)ZstdDecompressionWriter_write, METH_VARARGS,
+	{ "write", (PyCFunction)ZstdDecompressionWriter_write, METH_VARARGS | METH_KEYWORDS,
 	PyDoc_STR("Compress data") },
 	{ NULL, NULL }
 };