diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py --- a/Lib/test/test_zlib.py +++ b/Lib/test/test_zlib.py @@ -446,6 +446,13 @@ class CompressObjectTestCase(BaseCompres y = dco.decompress(x[:-5]) y += dco.flush() self.assertEqual(y, b'foo') + # flush(strict=True) raises an error, though + dco = zlib.decompressobj() + dco.decompress(x[:-5]) + with self.assertRaises(zlib.error) as cm: + dco.flush(strict=True) + self.assertEqual(str(cm.exception), + "Error -5 while decompressing data: incomplete or truncated stream") if hasattr(zlib.compressobj(), "copy"): def test_compresscopy(self): diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c --- a/Modules/zlibmodule.c +++ b/Modules/zlibmodule.c @@ -838,13 +838,16 @@ PyDoc_STRVAR(decomp_flush__doc__, "The decompressor object can no longer be used after this call."); static PyObject * -PyZlib_unflush(compobject *self, PyObject *args) +PyZlib_unflush(compobject *self, PyObject *args, PyObject *kwargs) { + static char *kwlist[] = {"length", "strict", NULL}; int err, length = DEFAULTALLOC; PyObject * retval = NULL; unsigned long start_total_out; + int strict = 0; - if (!PyArg_ParseTuple(args, "|i:flush", &length)) + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|ii:flush", + kwlist, &length, &strict)) return NULL; if (length <= 0) { PyErr_SetString(PyExc_ValueError, "length must be greater than zero"); @@ -866,7 +869,7 @@ PyZlib_unflush(compobject *self, PyObjec /* while Z_OK and the output buffer is full, there might be more output, so extend the output buffer and try again */ - while ((err == Z_OK || err == Z_BUF_ERROR) && self->zst.avail_out == 0) { + while (err == Z_OK || (err == Z_BUF_ERROR && self->zst.avail_out == 0)) { if (_PyBytes_Resize(&retval, length << 1) < 0) { Py_DECREF(retval); retval = NULL; @@ -894,6 +897,12 @@ PyZlib_unflush(compobject *self, PyObjec goto error; } } + else if (strict) { + zlib_error(self->zst, err, "while decompressing data"); + Py_DECREF(retval); + retval = NULL; + goto error; + } if (_PyBytes_Resize(&retval, self->zst.total_out - start_total_out) < 0) { Py_DECREF(retval); retval = NULL; @@ -923,7 +932,7 @@ static PyMethodDef Decomp_methods[] = { {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS, decomp_decompress__doc__}, - {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS, + {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS | METH_KEYWORDS, decomp_flush__doc__}, #ifdef HAVE_ZLIB_COPY {"copy", (PyCFunction)PyZlib_uncopy, METH_NOARGS,