[pypy-commit] pypy zlib-copying: Teach rzlib to copy decompression streams.

[Julian Berman]

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

Log:	Teach rzlib to copy decompression streams.

diff --git a/rpython/rlib/rzlib.py b/rpython/rlib/rzlib.py
--- a/rpython/rlib/rzlib.py
+++ b/rpython/rlib/rzlib.py
@@ -35,6 +35,7 @@
     MAX_WBITS  MAX_MEM_LEVEL
     Z_BEST_SPEED  Z_BEST_COMPRESSION  Z_DEFAULT_COMPRESSION
     Z_FILTERED  Z_HUFFMAN_ONLY  Z_DEFAULT_STRATEGY Z_NEED_DICT
+    Z_NULL
     '''.split()
 
 class SimpleCConfig:
@@ -160,6 +161,7 @@
     rffi.INT)
 _inflate = zlib_external('inflate', [z_stream_p, rffi.INT], rffi.INT)
 
+_inflateCopy = zlib_external('inflateCopy', [z_stream_p, z_stream_p], rffi.INT)
 _inflateEnd = zlib_external('inflateEnd', [z_stream_p], rffi.INT,
                             releasegil=False)
 
@@ -330,6 +332,19 @@
             lltype.free(stream, flavor='raw')
 
 
+def inflateCopy(source):
+    """
+    Allocate and return an independent copy of the provided stream object.
+    """
+    dest = inflateInit()
+    err = _inflateCopy(dest, source)
+    if err != Z_OK:
+        inflateEnd(dest)
+        raise RZlibError.fromstream(source, err,
+            "while copying decompression object")
+    return dest
+
+
 def inflateEnd(stream):
     """
     Free the resources associated with the inflate 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,60 @@
     rzlib.deflateEnd(stream)
 
 
+def test_decompress_copy():
+    """
+    inflateCopy produces an independent copy of a stream.
+    """
+
+    stream = rzlib.inflateInit()
+
+    bytes1, finished1, unused1 = rzlib.decompress(stream, compressed[:10])
+    assert bytes1
+    assert finished1 is False
+
+    copied = rzlib.inflateCopy(stream)
+
+    bytes_stream, finished_stream, unused_stream = rzlib.decompress(
+        stream,
+        compressed[10:],
+        rzlib.Z_FINISH,
+    )
+    assert bytes1 + bytes_stream == expanded
+    assert finished_stream is True
+    assert unused1 == 0
+    assert unused_stream == 0
+    rzlib.inflateEnd(stream)
+
+    bytes_copy, finished_copy, unused_copy = rzlib.decompress(
+        copied,
+        compressed[10:],
+        rzlib.Z_FINISH,
+    )
+    rzlib.inflateEnd(copied)
+    assert bytes1 + bytes_copy == expanded
+    assert finished_copy is True
+    assert unused_copy == 0
+
+
+def test_unsuccessful_decompress_copy():
+    """
+    Errors during unsuccesful inflateCopy operations raise RZlibErrors.
+    """
+    stream = rzlib.inflateInit()
+
+    # From zlib.h:
+    #
+    # "inflateCopy 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.inflateCopy, stream)
+    msg = "Error -2 while copying decompression object: inconsistent stream state"
+    assert str(exc.value) == msg
+    rzlib.inflateEnd(stream)
+
+
 def test_cornercases():
     """
     Test degenerate arguments.