Author: Julian Berman <julian...@grayvines.com>
Branch: 
Changeset: r95840:e0997a28d6e8
Date: 2019-02-05 15:12 +0100
http://bitbucket.org/pypy/pypy/changeset/e0997a28d6e8/

Log:    Add in the missing .copy methods for zlib (de)compression.

diff --git a/pypy/module/zlib/interp_zlib.py b/pypy/module/zlib/interp_zlib.py
--- a/pypy/module/zlib/interp_zlib.py
+++ b/pypy/module/zlib/interp_zlib.py
@@ -138,15 +138,18 @@
                  method=rzlib.Z_DEFLATED,             # \
                  wbits=rzlib.MAX_WBITS,               #  \   undocumented
                  memLevel=rzlib.DEF_MEM_LEVEL,        #  /    parameters
-                 strategy=rzlib.Z_DEFAULT_STRATEGY):  # /
+                 strategy=rzlib.Z_DEFAULT_STRATEGY,   # /
+                 stream=None):
         ZLibObject.__init__(self, space)
-        try:
-            self.stream = rzlib.deflateInit(level, method, wbits,
-                                            memLevel, strategy)
-        except rzlib.RZlibError as e:
-            raise zlib_error(space, e.msg)
-        except ValueError:
-            raise oefmt(space.w_ValueError, "Invalid initialization option")
+        if stream is None:
+            try:
+                stream = rzlib.deflateInit(level, method, wbits,
+                                                memLevel, strategy)
+            except rzlib.RZlibError as e:
+                raise zlib_error(space, e.msg)
+            except ValueError:
+                raise oefmt(space.w_ValueError, "Invalid initialization 
option")
+        self.stream = stream
         self.register_finalizer(space)
 
     def _finalize_(self):
@@ -178,6 +181,20 @@
             raise zlib_error(space, e.msg)
         return space.newbytes(result)
 
+    def copy(self, space):
+        """
+        copy() -- Return a copy of the compression object.
+        """
+        try:
+            self.lock()
+            try:
+                copied = rzlib.deflateCopy(self.stream)
+            finally:
+                self.unlock()
+        except rzlib.RZlibError as e:
+            raise zlib_error(space, e.msg)
+        return Compress(space=space, stream=copied)
+
     @unwrap_spec(mode="c_int")
     def flush(self, space, mode=rzlib.Z_FINISH):
         """
@@ -228,6 +245,7 @@
 Compress.typedef = TypeDef(
     'Compress',
     __new__ = interp2app(Compress___new__),
+    copy = interp2app(Compress.copy),
     compress = interp2app(Compress.compress),
     flush = interp2app(Compress.flush),
     __doc__ = """compressobj([level]) -- Return a compressor object.
@@ -241,7 +259,14 @@
     Wrapper around zlib's z_stream structure which provides convenient
     decompression functionality.
     """
-    def __init__(self, space, wbits=rzlib.MAX_WBITS):
+    def __init__(
+        self,
+        space,
+        wbits=rzlib.MAX_WBITS,
+        stream=None,
+        unused_data="",
+        unconsumed_tail="",
+    ):
         """
         Initialize a new decompression object.
 
@@ -251,14 +276,17 @@
         inflateInit2.
         """
         ZLibObject.__init__(self, space)
-        self.unused_data = ''
-        self.unconsumed_tail = ''
-        try:
-            self.stream = rzlib.inflateInit(wbits)
-        except rzlib.RZlibError as e:
-            raise zlib_error(space, e.msg)
-        except ValueError:
-            raise oefmt(space.w_ValueError, "Invalid initialization option")
+
+        if stream is None:
+            try:
+                stream = rzlib.inflateInit(wbits)
+            except rzlib.RZlibError as e:
+                raise zlib_error(space, e.msg)
+            except ValueError:
+                raise oefmt(space.w_ValueError, "Invalid initialization 
option")
+        self.stream = stream
+        self.unused_data = unused_data
+        self.unconsumed_tail = unconsumed_tail
         self.register_finalizer(space)
 
     def _finalize_(self):
@@ -305,6 +333,29 @@
         self._save_unconsumed_input(data, finished, unused_len)
         return space.newbytes(string)
 
+    def copy(self, space):
+        """
+        copy() -- Return a copy of the decompression object.
+        """
+        try:
+            self.lock()
+            try:
+                if not self.stream:
+                    raise zlib_error(space,
+                                     "decompressor object already flushed")
+                copied = rzlib.inflateCopy(self.stream)
+            finally:
+                self.unlock()
+        except rzlib.RZlibError as e:
+            raise zlib_error(space, e.msg)
+
+        return Decompress(
+            space=space,
+            stream=copied,
+            unused_data=self.unused_data,
+            unconsumed_tail=self.unconsumed_tail,
+        )
+
     def flush(self, space, w_length=None):
         """
         flush( [length] ) -- This is kept for backward compatibility,
@@ -345,6 +396,7 @@
 Decompress.typedef = TypeDef(
     'Decompress',
     __new__ = interp2app(Decompress___new__),
+    copy = interp2app(Decompress.copy),
     decompress = interp2app(Decompress.decompress),
     flush = interp2app(Decompress.flush),
     unused_data = interp_attrproperty('unused_data', Decompress, 
wrapfn="newbytes"),
diff --git a/pypy/module/zlib/test/test_zlib.py 
b/pypy/module/zlib/test/test_zlib.py
--- a/pypy/module/zlib/test/test_zlib.py
+++ b/pypy/module/zlib/test/test_zlib.py
@@ -8,8 +8,10 @@
 except ImportError:
     import py; py.test.skip("no zlib module on this host Python")
 
+from pypy.interpreter.gateway import interp2app
 try:
     from pypy.module.zlib import interp_zlib
+    from rpython.rlib import rzlib
 except ImportError:
     import py; py.test.skip("no zlib C library on this machine")
 
@@ -38,6 +40,22 @@
         cls.w_expanded = cls.space.wrap(expanded)
         cls.w_compressed = cls.space.wrap(zlib.compress(expanded))
 
+        def intentionally_break_a_z_stream(space, w_zobj):
+            """
+            Intentionally break the z_stream associated with a
+            compressobj or decompressobj in a way that causes their copy
+            methods to raise RZlibErrors.
+            """
+            from rpython.rtyper.lltypesystem import rffi, lltype
+            w_zobj.stream.c_zalloc = rffi.cast(
+                lltype.typeOf(w_zobj.stream.c_zalloc),
+                rzlib.Z_NULL,
+            )
+
+        cls.w_intentionally_break_a_z_stream = cls.space.wrap(
+            interp2app(intentionally_break_a_z_stream),
+        )
+
     def test_error(self):
         """
         zlib.error should be an exception class.
@@ -276,3 +294,54 @@
         assert dco.flush(1) == input1[1:]
         assert dco.unused_data == b''
         assert dco.unconsumed_tail == b''
+
+    def test_decompress_copy(self):
+        decompressor = self.zlib.decompressobj()
+        d1 = decompressor.decompress(self.compressed[:10])
+        assert d1
+
+        copied = decompressor.copy()
+
+        from_copy = copied.decompress(self.compressed[10:])
+        from_decompressor = decompressor.decompress(self.compressed[10:])
+
+        assert (d1 + from_copy) == (d1 + from_decompressor)
+
+    def test_unsuccessful_decompress_copy(self):
+        decompressor = self.zlib.decompressobj()
+        self.intentionally_break_a_z_stream(zobj=decompressor)
+        raises(self.zlib.error, decompressor.copy)
+
+    def test_decompress_copy_carries_along_state(self):
+        """
+        Decompressor.unused_data and unconsumed_tail are carried along when a
+        copy is done.
+        """
+        decompressor = self.zlib.decompressobj()
+        decompressor.decompress(self.compressed, max_length=5)
+        unconsumed_tail = decompressor.unconsumed_tail
+        assert unconsumed_tail
+        assert decompressor.copy().unconsumed_tail == unconsumed_tail
+
+        decompressor.decompress(decompressor.unconsumed_tail)
+        decompressor.decompress("xxx")
+        unused_data = decompressor.unused_data
+        assert unused_data
+        assert decompressor.copy().unused_data == unused_data
+
+    def test_compress_copy(self):
+        compressor = self.zlib.compressobj()
+        d1 = compressor.compress(self.expanded[:10])
+        assert d1
+
+        copied = compressor.copy()
+
+        from_copy = copied.compress(self.expanded[10:])
+        from_compressor = compressor.compress(self.expanded[10:])
+
+        assert (d1 + from_copy) == (d1 + from_compressor)
+
+    def test_unsuccessful_compress_copy(self):
+        compressor = self.zlib.compressobj()
+        self.intentionally_break_a_z_stream(zobj=compressor)
+        raises(self.zlib.error, compressor.copy)
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:
@@ -141,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)
 
@@ -160,6 +162,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)
 
@@ -290,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.
@@ -330,6 +346,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,108 @@
     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.
+    """
+
+    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.
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to