gh-160310: Share the one-shot decompress loop in compression._common

bz2.decompress(), lzma.decompress() and compression.zstd.decompress()
had three copies of the chunked multi-stream loop.  Move it into
compression._common._streams._decompress_all(), which takes the
decompressor factory, the trailing-data error, the stream magic and the
truncation error.  gzip keeps its own loop.

Author: gpshead

Lib/bz2.py

@@ -341,49 +341,11 @@ def decompress(data):
 
     For incremental decompression, use a BZ2Decompressor object instead.
     """
-    results = []
     if not data:
         return b""
-    # The views are released on exit, so that the caller can resize a
-    # bytearray argument again, also while handling an exception.
-    with memoryview(data) as view, view.cast("B") as data:
-        pos = 0
-        # There is usually a single stream: the first chunk is all of the
-        # data, to decompress it in one go.
-        chunk_size = len(data)
-        while True:
-            decomp = BZ2Decompressor()
-            start = pos
-            res = []
-            try:
-                # Feed the stream in chunks of increasing size, so that what
-                # is read past its end (and copied into unused_data) is
-                # bounded by the size of the stream, not by the size of all
-                # the streams that follow it.
-                while not decomp.eof and pos < len(data):
-                    with data[pos:pos + chunk_size] as chunk:
-                        res.append(decomp.decompress(chunk))
-                        pos += len(chunk)
-                    chunk_size *= 2
-            except OSError:
-                if not results:
-                    raise  # Error on the first iteration; bail out.
-                with data[start:start + 4] as head:
-                    if head.tobytes().startswith(_STREAM_MAGIC):
-                        # Leftover data is a damaged bzip2 stream; bail out.
-                        raise
-                break  # Leftover data is not a valid bzip2 stream; ignore it.
-            results.extend(res)
-            if not decomp.eof:
-                raise ValueError("Compressed data ended before the "
-                                 "end-of-stream marker was reached")
-            pos -= len(decomp.unused_data)
-            if pos == len(data):
-                break
-            # Later streams start small: a larger first chunk reads further
-            # past the end of a short stream, a smaller one needs more calls.
-            chunk_size = 1024
-    return b"".join(results)
+    return _streams._decompress_all(data, BZ2Decompressor, OSError,
+                                    stream_magic=_STREAM_MAGIC,
+                                    truncated_error=ValueError)
 
 
 def __getattr__(name):

Lib/compression/_common/_streams.py

@@ -169,3 +169,60 @@ def seek(self, offset, whence=io.SEEK_SET):
     def tell(self):
         """Return the current file position."""
         return self._pos
+
+
+def _decompress_all(data, decomp_factory, trailing_error=(), *,
+                    stream_magic=(), truncated_error, **decomp_args):
+    """Decompress all of the concatenated streams in *data*.
+
+    Each stream is decompressed by a new decompressor, made by calling
+    *decomp_factory* with *decomp_args*.  As in DecompressReader, an error
+    of class *trailing_error* in a stream after the first one ends the
+    data, unless that stream starts with *stream_magic*.  A stream without
+    an end-of-stream marker raises *truncated_error*.
+    """
+    results = []
+    # The views are released on exit, so that the caller can resize a
+    # bytearray argument again, also while handling an exception.
+    with memoryview(data) as view, view.cast("B") as data:
+        pos = 0
+        # There is usually a single stream: the first chunk is all of the
+        # data, to decompress it in one go.
+        chunk_size = len(data)
+        while True:
+            decomp = decomp_factory(**decomp_args)
+            start = pos
+            res = []
+            try:
+                # Feed the stream in chunks of increasing size, so that what
+                # is read past its end (and copied into unused_data) is
+                # bounded by the size of the stream, not by the size of all
+                # the streams that follow it.
+                while not decomp.eof and pos < len(data):
+                    with data[pos:pos + chunk_size] as chunk:
+                        res.append(decomp.decompress(chunk))
+                        pos += len(chunk)
+                    chunk_size *= 2
+            except trailing_error:
+                if not results:
+                    raise  # Error in the first stream; bail out.
+                if stream_magic:
+                    if isinstance(stream_magic, bytes):
+                        stream_magic = (stream_magic,)
+                    size = max(map(len, stream_magic))
+                    with data[start:start + size] as head:
+                        if head.tobytes().startswith(stream_magic):
+                            raise  # Damaged stream, not trailing data.
+                # Trailing data isn't a valid compressed stream; ignore it.
+                break
+            results.extend(res)
+            if not decomp.eof:
+                raise truncated_error("Compressed data ended before the "
+                                      "end-of-stream marker was reached")
+            pos -= len(decomp.unused_data)
+            if pos == len(data):
+                break
+            # Later streams start small: a larger first chunk reads further
+            # past the end of a short stream, a smaller one needs more calls.
+            chunk_size = 1024
+    return b"".join(results)

Lib/compression/zstd/__init__.py

@@ -33,6 +33,7 @@
 from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdDict, ZstdError,
                    get_frame_size, zstd_version, ZSTD_VERSION,
                    zstd_version_info, ZSTD_VERSION_INFO)
+from compression._common import _streams
 from compression.zstd._zstdfile import ZstdFile, open, _nbytes
 
 
@@ -160,35 +161,9 @@ def decompress(data, zstd_dict=None, options=None):
 
     For incremental decompression, use a ZstdDecompressor instead.
     """
-    results = []
-    # The views are released on exit, so that the caller can resize a
-    # bytearray argument again, also while handling an exception.
-    with memoryview(data) as view, view.cast('B') as data:
-        pos = 0
-        # There is usually a single frame: the first chunk is all of the
-        # data, to decompress it in one go.
-        chunk_size = len(data)
-        while True:
-            decomp = ZstdDecompressor(options=options, zstd_dict=zstd_dict)
-            # Feed the frame in chunks of increasing size, so that what is
-            # read past its end (and copied into unused_data) is bounded by
-            # the size of the frame, not by the size of all the frames that
-            # follow it.
-            while not decomp.eof and pos < len(data):
-                with data[pos:pos + chunk_size] as chunk:
-                    results.append(decomp.decompress(chunk))
-                    pos += len(chunk)
-                chunk_size *= 2
-            if not decomp.eof:
-                raise ZstdError('Compressed data ended before the '
-                                'end-of-stream marker was reached')
-            pos -= len(decomp.unused_data)
-            if pos == len(data):
-                break
-            # Later frames start small: a larger first chunk reads further
-            # past the end of a short frame, a smaller one needs more calls.
-            chunk_size = 1024
-    return b''.join(results)
+    return _streams._decompress_all(data, ZstdDecompressor,
+                                    truncated_error=ZstdError,
+                                    options=options, zstd_dict=zstd_dict)
 
 
 class CompressionParameter(enum.IntEnum):

Lib/lzma.py

@@ -354,45 +354,8 @@ def decompress(data, format=FORMAT_AUTO, memlimit=None, filters=None):
 
     For incremental decompression, use an LZMADecompressor instead.
     """
-    results = []
-    magic = _stream_magic(format)
-    # The views are released on exit, so that the caller can resize a
-    # bytearray argument again, also while handling an exception.
-    with memoryview(data) as view, view.cast("B") as data:
-        pos = 0
-        # There is usually a single stream: the first chunk is all of the
-        # data, to decompress it in one go.
-        chunk_size = len(data)
-        while True:
-            decomp = LZMADecompressor(format, memlimit, filters)
-            start = pos
-            res = []
-            try:
-                # Feed the stream in chunks of increasing size, so that what
-                # is read past its end (and copied into unused_data) is
-                # bounded by the size of the stream, not by the size of all
-                # the streams that follow it.
-                while not decomp.eof and pos < len(data):
-                    with data[pos:pos + chunk_size] as chunk:
-                        res.append(decomp.decompress(chunk))
-                        pos += len(chunk)
-                    chunk_size *= 2
-            except LZMAError:
-                if not results:
-                    raise  # Error on the first iteration; bail out.
-                with data[start:start + len(magic)] as head:
-                    if magic and head.tobytes().startswith(magic):
-                        # Leftover data is a damaged XZ stream; bail out.
-                        raise
-                break  # Leftover data is not a valid LZMA/XZ stream; ignore it.
-            results.extend(res)
-            if not decomp.eof:
-                raise LZMAError("Compressed data ended before the "
-                                "end-of-stream marker was reached")
-            pos -= len(decomp.unused_data)
-            if pos == len(data):
-                break
-            # Later streams start small: a larger first chunk reads further
-            # past the end of a short stream, a smaller one needs more calls.
-            chunk_size = 1024
-    return b"".join(results)
+    return _streams._decompress_all(data, LZMADecompressor, LZMAError,
+                                    stream_magic=_stream_magic(format),
+                                    truncated_error=LZMAError,
+                                    format=format, memlimit=memlimit,
+                                    filters=filters)

Lib/test/test_bz2.py

@@ -1126,6 +1126,10 @@ def testDecompressMultiStreamBig(self):
         self.assertEqual(bz2.decompress(memoryview(data)), expected)
         self.assertEqual(bz2.decompress(data + self.BAD_DATA), expected)
         self.assertRaises(ValueError, bz2.decompress, data[:-50_000])
+        # The damage is past the first chunk fed for the second stream.
+        damaged = bytearray(bz2.compress(big))
+        damaged[-3] ^= 1
+        self.assertRaises(OSError, bz2.decompress, self.DATA + damaged)
 
     def testDecompressErrorReleasesBuffer(self):
         # The argument can be resized while the exception is handled.

Lib/test/test_lzma.py

@@ -524,6 +524,10 @@ def test_decompress_multistream_big(self):
         self.assertEqual(lzma.decompress(memoryview(cdata)), expected)
         self.assertEqual(lzma.decompress(cdata + COMPRESSED_BOGUS), expected)
         self.assertRaises(LZMAError, lzma.decompress, cdata[:-50_000])
+        # The damage is past the first chunk fed for the second stream.
+        damaged = bytearray(lzma.compress(big))
+        damaged[-1] ^= 1
+        self.assertRaises(LZMAError, lzma.decompress, COMPRESSED_XZ + damaged)
 
     def test_decompress_error_releases_buffer(self):
         # The argument can be resized while the exception is handled.
@@ -563,6 +567,9 @@ def decompress(self, data):
     def test_decompress_trailing_junk(self):
         ddata = lzma.decompress(COMPRESSED_XZ + COMPRESSED_BOGUS)
         self.assertEqual(ddata, INPUT)
+        ddata = lzma.decompress(COMPRESSED_ALONE + COMPRESSED_BOGUS,
+                                lzma.FORMAT_ALONE)
+        self.assertEqual(ddata, INPUT)
 
     def test_decompress_multistream_trailing_junk(self):
         ddata = lzma.decompress(COMPRESSED_XZ * 3 + COMPRESSED_BOGUS)