decompressionwriter: use buffer protocol in write() (#26)

c-ext/decompressionwriter.c

@@ -55,18 +55,17 @@ static PyObject* ZstdDecompressionWriter_memory_size(ZstdDecompressionWriter* se
 
 static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, PyObject* args) {
 	PyObject* result = NULL;
-	const char* source;
-	Py_ssize_t sourceSize;
+	Py_buffer source;
 	size_t zresult = 0;
 	ZSTD_inBuffer input;
 	ZSTD_outBuffer output;
 	PyObject* res;
 	Py_ssize_t totalWrite = 0;
 
 #if PY_MAJOR_VERSION >= 3
-	if (!PyArg_ParseTuple(args, "y#:write", &source, &sourceSize)) {
+	if (!PyArg_ParseTuple(args, "y*:write", &source)) {
 #else
-	if (!PyArg_ParseTuple(args, "s#:write", &source, &sourceSize)) {
+	if (!PyArg_ParseTuple(args, "s*:write", &source)) {
 #endif
 		return NULL;
 	}
@@ -86,11 +85,11 @@ static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, Py
 	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);
 		Py_END_ALLOW_THREADS
@@ -120,6 +119,7 @@ static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, Py
 	result = PyLong_FromSsize_t(totalWrite);
 
 finally:
+	PyBuffer_Release(&source);
 	return result;
 }
 

tests/test_decompressor.py

@@ -386,6 +386,23 @@ def test_empty_roundtrip(self):
         empty = cctx.compress(b'')
         self.assertEqual(decompress_via_writer(empty), b'')
 
+    def test_input_types(self):
+        cctx = zstd.ZstdCompressor(level=1)
+        compressed = cctx.compress(b'foo')
+
+        sources = [
+            memoryview(compressed),
+            bytearray(compressed),
+        ]
+
+        dctx = zstd.ZstdDecompressor()
+        for source in sources:
+            buffer = io.BytesIO()
+            with dctx.write_to(buffer) as decompressor:
+                decompressor.write(source)
+
+            self.assertEqual(buffer.getvalue(), b'foo')
+
     def test_large_roundtrip(self):
         chunks = []
         for i in range(255):