Add decompression

try.py

@@ -9,3 +9,8 @@ print("res", res)
 print(b[:res])
 
 print("Roundtrip", zlib.decompress(b[:res]) == a)
+
+c = bytearray(25001)
+res2 = zlib_into.decompress_into(b[:res], c)
+print(f"Decompressed {res2} bytes")
+print("Roundtrip within zlib_into", bytes(c[:res2]) == a)

zlib_into.c

@@ -99,7 +99,7 @@ compress_into(PyObject *module, PyObject *args, PyObject *kwargs) {
         goto done;
     default:
         deflateEnd(&zst);
-        zlib_error(zst, err, "while compressing data");
+        zlib_error(zst, err, "while preparing to compress data");
         goto done;
     }
 
@@ -121,7 +121,11 @@ compress_into(PyObject *module, PyObject *args, PyObject *kwargs) {
             goto done;
     }
     bytes_written = output.len - zst.avail_out;
-    deflateEnd(&zst);
+    err = deflateEnd(&zst);
+    if (err != Z_OK) {
+        zlib_error(zst, err, "while finishing compression");
+        goto done;
+    }
 
     return_value = PyLong_FromSsize_t(bytes_written);
 
@@ -129,11 +133,87 @@ compress_into(PyObject *module, PyObject *args, PyObject *kwargs) {
       PyBuffer_Release(&input);
       PyBuffer_Release(&output);
       return return_value;
+}
 
+
+static PyObject *
+decompress_into(PyObject *module, PyObject *args, PyObject *kwargs) {
+    static char *keywords[] = {"data", "output", "wbits", NULL};
+    PyObject *return_value = NULL;
+    Py_ssize_t bytes_written;
+    Py_buffer input, output;
+    int wbits=MAX_WBITS;
+    z_stream zst;
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*w*|i", keywords,
+                                     &input, &output, &wbits)) {
+        return NULL;
+    }
+
+    if (output.len <= 0) {
+        PyErr_SetString(PyExc_ValueError, "Output buffer may not be 0 size");
+        goto done;
+    }
+
+    zst.opaque = NULL;
+    zst.zalloc = PyZlib_Malloc;
+    zst.zfree = PyZlib_Free;
+    zst.next_in = input.buf;
+    zst.avail_in = input.len;
+    zst.next_out = output.buf;
+    zst.avail_out = output.len;
+    int err = inflateInit2(&zst, wbits);
+
+    switch (err) {
+    case Z_OK:
+        break;
+    case Z_MEM_ERROR:
+        PyErr_SetString(PyExc_MemoryError,
+                        "Out of memory while decompressing data");
+        goto done;
+    default:
+        deflateEnd(&zst);
+        zlib_error(zst, err, "while preparing to decompress data");
+        goto done;
+    }
+
+    Py_BEGIN_ALLOW_THREADS
+    err = inflate(&zst, Z_FINISH);
+    Py_END_ALLOW_THREADS
+
+    switch (err) {
+        case Z_STREAM_END:
+            break;
+        case Z_OK:
+        case Z_BUF_ERROR:
+            deflateEnd(&zst);
+            PyErr_SetString(PyExc_BufferError, "Not enough space in output buffer");
+            goto done;
+        default:
+            deflateEnd(&zst);
+            zlib_error(zst, err, "while compressing data");
+            goto done;
+    }
+    bytes_written = output.len - zst.avail_out;
+
+    err = inflateEnd(&zst);
+    if (err != Z_OK) {
+        zlib_error(zst, err, "while finishing decompression");
+        goto done;
+    }
+
+    return_value = PyLong_FromSsize_t(bytes_written);
+
+    done:
+      PyBuffer_Release(&input);
+      PyBuffer_Release(&output);
+      return return_value;
 }
 
+
 static PyMethodDef ZlibIntoMethods[] = {
     {"compress_into", (PyCFunction)compress_into, METH_VARARGS | METH_KEYWORDS, "zlib compress data into a buffer"},
+    {"decompress_into", (PyCFunction)decompress_into, METH_VARARGS | METH_KEYWORDS, "zlib decompress data into a buffer"},
     {NULL, NULL, 0, NULL}
 };