Patch #1435422: zlib's compress and decompress objects now have a

copy() method.

4 files changed, 175 insertions, 0 deletions

diff --git a/Doc/lib/libzlib.tex b/Doc/lib/libzlib.tex
index dfbb43d..876f8c0 100644
--- a/Doc/lib/libzlib.tex
+++ b/Doc/lib/libzlib.tex
@@ -123,6 +123,12 @@ prevents compressing any more data.  After calling
 action is to delete the object.  
 \end{methoddesc}
 
+\begin{methoddesc}[Compress]{copy}{}
+Returns a copy of the compression object.  This can be used to efficiently
+compress a set of data that share a common initial prefix.
+\versionadded{2.5}
+\end{methoddesc}
+
 Decompression objects support the following methods, and two attributes:
 
 \begin{memberdesc}{unused_data}
@@ -176,6 +182,13 @@ The optional parameter \var{length} sets the initial size of the
 output buffer.
 \end{methoddesc}
 
+\begin{methoddesc}[Decompress]{copy}{}
+Returns a copy of the decompression object.  This can be used to save the
+state of the decompressor midway through the data stream in order to speed up
+random seeks into the stream at a future point.
+\versionadded{2.5}
+\end{methoddesc}
+
 \begin{seealso}
   \seemodule{gzip}{Reading and writing \program{gzip}-format files.}
   \seeurl{http://www.zlib.net}{The zlib library home page.}
diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index 7634680..ccbc8fd 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -302,6 +302,63 @@ class CompressObjectTestCase(unittest.TestCase):
         dco = zlib.decompressobj()
         self.assertEqual(dco.flush(), "") # Returns nothing
 
+    def test_compresscopy(self):
+        # Test copying a compression object
+        data0 = HAMLET_SCENE
+        data1 = HAMLET_SCENE.swapcase()
+        c0 = zlib.compressobj(zlib.Z_BEST_COMPRESSION)
+        bufs0 = []
+        bufs0.append(c0.compress(data0))
+
+        c1 = c0.copy()
+        bufs1 = bufs0[:]
+
+        bufs0.append(c0.compress(data0))
+        bufs0.append(c0.flush())
+        s0 = ''.join(bufs0)
+
+        bufs1.append(c1.compress(data1))
+        bufs1.append(c1.flush())
+        s1 = ''.join(bufs1)
+
+        self.assertEqual(zlib.decompress(s0),data0+data0)
+        self.assertEqual(zlib.decompress(s1),data0+data1)
+
+    def test_badcompresscopy(self):
+        # Test copying a compression object in an inconsistent state
+        c = zlib.compressobj()
+        c.compress(HAMLET_SCENE)
+        c.flush()
+        self.assertRaises(ValueError, c.copy)
+
+    def test_decompresscopy(self):
+        # Test copying a decompression object
+        data = HAMLET_SCENE
+        comp = zlib.compress(data)
+
+        d0 = zlib.decompressobj()
+        bufs0 = []
+        bufs0.append(d0.decompress(comp[:32]))
+
+        d1 = d0.copy()
+        bufs1 = bufs0[:]
+
+        bufs0.append(d0.decompress(comp[32:]))
+        s0 = ''.join(bufs0)
+
+        bufs1.append(d1.decompress(comp[32:]))
+        s1 = ''.join(bufs1)
+
+        self.assertEqual(s0,s1)
+        self.assertEqual(s0,data)
+
+    def test_baddecompresscopy(self):
+        # Test copying a compression object in an inconsistent state
+        data = zlib.compress(HAMLET_SCENE)
+        d = zlib.decompressobj()
+        d.decompress(data)
+        d.flush()
+        self.assertRaises(ValueError, d.copy)
 
 def genblock(seed, length, step=1024, generator=random):
     """length-byte stream of random data from a seed (in step-byte blocks)."""
diff --git a/Misc/NEWS b/Misc/NEWS
index 6551b99..20d4ff1 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -28,6 +28,9 @@ Core and builtins
 Extension Modules
 -----------------
 
+- Patch #1435422: zlib's compress and decompress objects now have a
+  copy() method.
+
 - On Win32, os.listdir now supports arbitrarily-long Unicode path names
   (up to the system limit of 32K characters).
 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 35b8c32..b41ed50 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -653,6 +653,104 @@ PyZlib_flush(compobject *self, PyObject *args)
     return RetVal;
 }
 
+PyDoc_STRVAR(comp_copy__doc__,
+"copy() -- Return a copy of the compression object.");
+
+static PyObject *
+PyZlib_copy(compobject *self)
+{
+    compobject *retval = NULL;
+    int err;
+
+    retval = newcompobject(&Comptype);
+    if (!retval) return NULL;
+
+    /* Copy the zstream state
+     * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
+     */
+    ENTER_ZLIB
+    err = deflateCopy(&retval->zst, &self->zst);
+    switch(err) {
+    case(Z_OK):
+        break;
+    case(Z_STREAM_ERROR):
+        PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
+        goto error;
+    case(Z_MEM_ERROR):
+        PyErr_SetString(PyExc_MemoryError,
+                        "Can't allocate memory for compression object");
+        goto error;
+    default:
+        zlib_error(self->zst, err, "while copying compression object");
+        goto error;
+    }
+
+    retval->unused_data = self->unused_data;
+    retval->unconsumed_tail = self->unconsumed_tail;
+    Py_INCREF(retval->unused_data);
+    Py_INCREF(retval->unconsumed_tail);
+
+    /* Mark it as being initialized */
+    retval->is_initialised = 1;
+
+    LEAVE_ZLIB
+    return (PyObject *)retval;
+
+error:
+    LEAVE_ZLIB
+    Py_XDECREF(retval);	
+    return NULL;
+}
+
+PyDoc_STRVAR(decomp_copy__doc__,
+"copy() -- Return a copy of the decompression object.");
+
+static PyObject *
+PyZlib_uncopy(compobject *self)
+{
+    compobject *retval = NULL;
+    int err;
+
+    retval = newcompobject(&Decomptype);
+    if (!retval) return NULL;
+
+    /* Copy the zstream state
+     * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
+     */
+    ENTER_ZLIB
+    err = inflateCopy(&retval->zst, &self->zst);
+    switch(err) {
+    case(Z_OK):
+        break;
+    case(Z_STREAM_ERROR):
+        PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
+        goto error;
+    case(Z_MEM_ERROR):
+        PyErr_SetString(PyExc_MemoryError,
+                        "Can't allocate memory for decompression object");
+        goto error;
+    default:
+        zlib_error(self->zst, err, "while copying decompression object");
+        goto error;
+    }
+
+    retval->unused_data = self->unused_data;
+    retval->unconsumed_tail = self->unconsumed_tail;
+    Py_INCREF(retval->unused_data);
+    Py_INCREF(retval->unconsumed_tail);
+
+    /* Mark it as being initialized */
+    retval->is_initialised = 1;
+
+    LEAVE_ZLIB
+    return (PyObject *)retval;
+
+error:
+    LEAVE_ZLIB
+    Py_XDECREF(retval);
+    return NULL;
+}
+
 PyDoc_STRVAR(decomp_flush__doc__,
 "flush( [length] ) -- Return a string containing any remaining\n"
 "decompressed data. length, if given, is the initial size of the\n"
@@ -725,6 +823,8 @@ static PyMethodDef comp_methods[] =
                  comp_compress__doc__},
     {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
               comp_flush__doc__},
+    {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
+              comp_copy__doc__},
     {NULL, NULL}
 };
 
@@ -734,6 +834,8 @@ static PyMethodDef Decomp_methods[] =
                    decomp_decompress__doc__},
     {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
               decomp_flush__doc__},
+    {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
+              decomp_copy__doc__},
     {NULL, NULL}
 };