diff -r 2f563908ebc5 Lib/test/test_zlib.py --- a/Lib/test/test_zlib.py Thu Apr 26 17:05:31 2012 +0200 +++ b/Lib/test/test_zlib.py Thu May 03 17:50:36 2012 -0700 @@ -425,6 +425,35 @@ dco = zlib.decompressobj() self.assertEqual(dco.flush(), b"") # Returns nothing + def test_dictionary(self): + # test short string + zdict = b"thequickbrownfoxjumped\x00" + co = zlib.compressobj(dict=zdict) + x = b"the quick brown fox jumped over the candlestick" + cd = co.compress(x) + co.flush() + dco = zlib.decompressobj(dict=zdict) + self.assertEqual((dco.decompress (cd) + dco.flush()), x) + # try again without a dict + dco = zlib.decompressobj() + self.assertRaises(zlib.error, dco.decompress, cd) + # take a chunk out of the middle of hamlet + co = zlib.compressobj(dict=HAMLET_SCENE) + lh3 = len(HAMLET_SCENE) // 3 + chunk = HAMLET_SCENE[lh3:lh3*2] + cd = co.compress (chunk) + co.flush() + dco = zlib.decompressobj(dict=HAMLET_SCENE) + self.assertEqual (dco.decompress (cd) + dco.flush(), chunk) + # test streaming behavior with dictionary + co = zlib.compressobj (dict=HAMLET_SCENE) + do = zlib.decompressobj (dict=HAMLET_SCENE) + piece = HAMLET_SCENE[1000:1500] + d0 = co.compress (piece) + co.flush (zlib.Z_SYNC_FLUSH) + d1 = co.compress (piece[100:]) + co.flush (zlib.Z_SYNC_FLUSH) + d2 = co.compress (piece[:-100]) + co.flush (zlib.Z_SYNC_FLUSH) + self.assertEqual (do.decompress (d0), piece) + self.assertEqual (do.decompress (d1), piece[100:]) + self.assertEqual (do.decompress (d2), piece[:-100]) + def test_decompress_incomplete_stream(self): # This is 'foo', deflated x = b'x\x9cK\xcb\xcf\x07\x00\x02\x82\x01E' diff -r 2f563908ebc5 Modules/zlibmodule.c --- a/Modules/zlibmodule.c Thu Apr 26 17:05:31 2012 +0200 +++ b/Modules/zlibmodule.c Thu May 03 17:50:36 2012 -0700 @@ -45,6 +45,7 @@ PyObject *unconsumed_tail; char eof; int is_initialised; + PyObject *dict; #ifdef WITH_THREAD PyThread_type_lock lock; #endif @@ -98,6 +99,7 @@ return NULL; self->eof = 0; self->is_initialised = 0; + self->dict = NULL; self->unused_data = PyBytes_FromStringAndSize("", 0); if (self->unused_data == NULL) { Py_DECREF(self); @@ -316,14 +318,17 @@ } static PyObject * -PyZlib_compressobj(PyObject *selfptr, PyObject *args) +PyZlib_compressobj(PyObject *selfptr, PyObject *args, PyObject *kwargs) { compobject *self; int level=Z_DEFAULT_COMPRESSION, method=DEFLATED; int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err; + PyObject * dict=NULL; + static char *kwlist[] = {"level", "method", "wbits", "memLevel", "strategy", "dict", NULL}; - if (!PyArg_ParseTuple(args, "|iiiii:compressobj", &level, &method, &wbits, - &memLevel, &strategy)) + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iiiiiO!:compressobj", + kwlist, &level, &method, &wbits, + &memLevel, &strategy, &PyBytes_Type, &dict)) return NULL; self = newcompobject(&Comptype); @@ -337,7 +342,27 @@ switch(err) { case (Z_OK): self->is_initialised = 1; - return (PyObject*)self; + if (!dict) { + return (PyObject*)self; + } else { + err = deflateSetDictionary (&self->zst, + (Byte*)PyBytes_AsString (dict), + PyBytes_Size(dict)); + switch (err) { + case (Z_OK): + Py_INCREF(dict); + self->dict = dict; + return (PyObject*)self; + case (Z_STREAM_ERROR): + Py_DECREF(self); + PyErr_SetString(PyExc_ValueError, "Invalid dictionary"); + return NULL; + default: + Py_DECREF(self); + PyErr_SetString(PyExc_ValueError, "deflateSetDictionary()"); + return NULL; + } + } case (Z_MEM_ERROR): Py_DECREF(self); PyErr_SetString(PyExc_MemoryError, @@ -355,11 +380,14 @@ } static PyObject * -PyZlib_decompressobj(PyObject *selfptr, PyObject *args) +PyZlib_decompressobj(PyObject *selfptr, PyObject *args, PyObject *kwargs) { + static char *kwlist[] = {"wbits", "dict", NULL}; int wbits=DEF_WBITS, err; compobject *self; - if (!PyArg_ParseTuple(args, "|i:decompressobj", &wbits)) + PyObject *dict=NULL; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!:decompressobj", + kwlist, &wbits, &PyBytes_Type, &dict)) return NULL; self = newcompobject(&Decomptype); @@ -369,6 +397,10 @@ self->zst.zfree = (free_func)Z_NULL; self->zst.next_in = NULL; self->zst.avail_in = 0; + if (dict) { + Py_INCREF(dict); + self->dict = dict; + } err = inflateInit2(&self->zst, wbits); switch(err) { case (Z_OK): @@ -398,6 +430,7 @@ #endif Py_XDECREF(self->unused_data); Py_XDECREF(self->unconsumed_tail); + Py_XDECREF(self->dict); PyObject_Del(self); } @@ -557,6 +590,20 @@ err = inflate(&(self->zst), Z_SYNC_FLUSH); Py_END_ALLOW_THREADS + if (err == Z_NEED_DICT && self->dict) { + err = inflateSetDictionary(&(self->zst), + (Byte*)PyBytes_AsString (self->dict), + PyBytes_Size (self->dict)); + if (err != Z_OK) { + zlib_error(self->zst, err, "while decompressing data"); + Py_DECREF(RetVal); + RetVal = NULL; + goto error; + } + /* repeat the call to inflate! */ + err = inflate(&(self->zst), Z_SYNC_FLUSH); + } + /* While Z_OK and the output buffer is full, there might be more output. So extend the output buffer and try again. */ @@ -770,10 +817,13 @@ } Py_INCREF(self->unused_data); Py_INCREF(self->unconsumed_tail); + Py_XINCREF(self->dict); Py_XDECREF(retval->unused_data); Py_XDECREF(retval->unconsumed_tail); + Py_XDECREF(retval->dict); retval->unused_data = self->unused_data; retval->unconsumed_tail = self->unconsumed_tail; + retval->dict = self->dict; retval->eof = self->eof; /* Mark it as being initialized */ @@ -822,10 +872,13 @@ Py_INCREF(self->unused_data); Py_INCREF(self->unconsumed_tail); + Py_XINCREF(self->dict); Py_XDECREF(retval->unused_data); Py_XDECREF(retval->unconsumed_tail); + Py_XDECREF(retval->dict); retval->unused_data = self->unused_data; retval->unconsumed_tail = self->unconsumed_tail; + retval->dict = self->dict; retval->eof = self->eof; /* Mark it as being initialized */ @@ -1032,13 +1085,13 @@ adler32__doc__}, {"compress", (PyCFunction)PyZlib_compress, METH_VARARGS, compress__doc__}, - {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS, + {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS|METH_KEYWORDS, compressobj__doc__}, {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS, crc32__doc__}, {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS, decompress__doc__}, - {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS, + {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS|METH_KEYWORDS, decompressobj__doc__}, {NULL, NULL} };