[pypy-commit] pypy zlib-copying: And the same for compression streams.

[Julian Berman]

Author: Julian Berman <Julian+Hg at GrayVines.com>
Branch: zlib-copying
Changeset: r95834:17e6b069ffdd
Date: 2019-02-04 17:21 +0100
http://bitbucket.org/pypy/pypy/changeset/17e6b069ffdd/

Log:	And the same for compression streams.

diff --git a/rpython/rlib/rzlib.py b/rpython/rlib/rzlib.py
--- a/rpython/rlib/rzlib.py
+++ b/rpython/rlib/rzlib.py
@@ -142,6 +142,7 @@
     rffi.INT)
 _deflate = zlib_external('deflate', [z_stream_p, rffi.INT], rffi.INT)
 
+_deflateCopy = zlib_external('deflateCopy', [z_stream_p, z_stream_p], rffi.INT)
 _deflateEnd = zlib_external('deflateEnd', [z_stream_p], rffi.INT,
                             releasegil=False)
 
@@ -292,6 +293,19 @@
             lltype.free(stream, flavor='raw')
 
 
+def deflateCopy(source):
+    """
+    Allocate and return an independent copy of the provided stream object.
+    """
+    dest = deflateInit()
+    err = _deflateCopy(dest, source)
+    if err != Z_OK:
+        deflateEnd(dest)
+        raise RZlibError.fromstream(source, err,
+            "while copying compression object")
+    return dest
+
+
 def deflateEnd(stream):
     """
     Free the resources associated with the deflate stream.
diff --git a/rpython/rlib/test/test_rzlib.py b/rpython/rlib/test/test_rzlib.py
--- a/rpython/rlib/test/test_rzlib.py
+++ b/rpython/rlib/test/test_rzlib.py
@@ -246,6 +246,54 @@
     rzlib.deflateEnd(stream)
 
 
+def test_compress_copy():
+    """
+    inflateCopy produces an independent copy of a stream.
+    """
+
+    stream = rzlib.deflateInit()
+
+    bytes1 = rzlib.compress(stream, expanded[:10])
+    assert bytes1
+
+    copied = rzlib.deflateCopy(stream)
+
+    bytes_stream = rzlib.compress(
+        stream,
+        expanded[10:],
+        rzlib.Z_FINISH,
+    )
+    assert bytes1 + bytes_stream == compressed
+    rzlib.deflateEnd(stream)
+
+    bytes_copy = rzlib.compress(
+        copied,
+        expanded[10:],
+        rzlib.Z_FINISH,
+    )
+    rzlib.deflateEnd(copied)
+    assert bytes1 + bytes_copy == compressed
+
+
+def test_unsuccessful_compress_copy():
+    """
+    Errors during unsuccesful deflateCopy operations raise RZlibErrors.
+    """
+    stream = rzlib.deflateInit()
+
+    # From zlib.h:
+    #
+    # "deflateCopy returns [...] Z_STREAM_ERROR if the source stream
+    #  state was inconsistent (such as zalloc being Z_NULL)"
+    from rpython.rtyper.lltypesystem import rffi, lltype
+    stream.c_zalloc = rffi.cast(lltype.typeOf(stream.c_zalloc), rzlib.Z_NULL)
+
+    exc = py.test.raises(rzlib.RZlibError, rzlib.deflateCopy, stream)
+    msg = "Error -2 while copying compression object: inconsistent stream state"
+    assert str(exc.value) == msg
+    rzlib.deflateEnd(stream)
+
+
 def test_decompress_copy():
     """
     inflateCopy produces an independent copy of a stream.