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