[pypy-commit] pypy zlib-copying: And now the same for compressobj...

[Julian Berman]

Author: Julian Berman <Julian+Hg at GrayVines.com>
Branch: zlib-copying
Changeset: r95837:9b6800b03333
Date: 2019-02-05 14:12 +0100
http://bitbucket.org/pypy/pypy/changeset/9b6800b03333/

Log:	And now the same for compressobj...

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,17 @@
             raise zlib_error(space, e.msg)
         return space.newbytes(result)
 
+    def copy(self, space):
+        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 +242,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.
@@ -266,7 +281,6 @@
                 raise zlib_error(space, e.msg)
             except ValueError:
                 raise oefmt(space.w_ValueError, "Invalid initialization option")
-        self.space = space
         self.stream = stream
         self.unused_data = unused_data
         self.unconsumed_tail = unconsumed_tail
@@ -330,7 +344,7 @@
             raise zlib_error(space, e.msg)
 
         return Decompress(
-            space=self.space,
+            space=space,
             stream=copied,
             # unused_data=self.unused_data,
             # unconsumed_tail=self.unconsumed_tail,
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
@@ -311,3 +311,20 @@
         decompressor = self.zlib.decompressobj()
         self.intentionally_break_a_z_stream(zobj=decompressor)
         raises(self.zlib.error, decompressor.copy)
+
+    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)