gh-155152: Improve performance of concatenated zstd frames in zipimport (#155154)

Co-authored-by: Emma Smith <emma@emmatyping.dev>

Author: lpyu001

Lib/test/test_zipimport.py

@@ -1102,6 +1102,62 @@ class DeflateCompressedZipImportTestCase(UncompressedZipImportTestCase):
 class ZStdCompressedZipImportTestCase(UncompressedZipImportTestCase):
     compression = ZIP_ZSTANDARD
 
+    def test_concatenated_frames(self):
+        from compression import zstd
+        import zipfile
+
+        class ConcatenatedFrameCompressor:
+            def compress(self, data):
+                return b"".join(zstd.compress(bytes([byte])) for byte in data)
+
+            def flush(self):
+                return b""
+
+        with unittest.mock.patch.object(
+            zipfile,
+            "_get_compressor",
+            return_value=ConcatenatedFrameCompressor(),
+        ):
+            self.doTest(".py", {TESTMOD + ".py": test_src}, TESTMOD)
+
+    def test_concatenated_frames_linear_input(self):
+        from compression import zstd
+
+        frame_count = 100
+        frame = zstd.compress(b"x")
+        compressed = frame * frame_count
+        input_sizes = []
+        decompressor_class = zipimport._get_zstd_decompressor_class()
+
+        class RecordingDecompressor:
+            def __init__(self):
+                self._decompressor = decompressor_class()
+
+            def __getattr__(self, name):
+                return getattr(self._decompressor, name)
+
+            def decompress(self, data, max_length=-1):
+                input_sizes.append(memoryview(data).nbytes)
+                return self._decompressor.decompress(data, max_length)
+
+        with unittest.mock.patch.object(
+            zipimport, "_zstd_decompressor_class", RecordingDecompressor
+        ):
+            self.assertEqual(
+                zipimport._zstd_decompress(compressed), b"x" * frame_count
+            )
+
+        self.assertEqual(sum(input_sizes), len(compressed))
+
+    def test_truncated_frame(self):
+        from compression import zstd
+
+        compressed = zstd.compress(b"x")
+        for data in (b"", compressed[:-1]):
+            with self.subTest(data=data):
+                with self.assertRaises(zipimport.ZipImportError):
+                    zipimport._zstd_decompress(data)
+
 
 class BadFileZipImportTestCase(unittest.TestCase):
     def assertZipFailure(self, filename):

Lib/zipimport.py

@@ -585,12 +585,13 @@ def _get_zlib_decompress_func():
 
 _importing_zstd = False
 _zstd_decompressor_class = None
+_zstd_get_frame_size = None
 
 # Return the _zstd.ZstdDecompressor function object, or NULL if _zstd couldn't
 # be imported. The result is cached when found.
 def _get_zstd_decompressor_class():
-    global _zstd_decompressor_class
-    if _zstd_decompressor_class:
+    global _zstd_decompressor_class, _zstd_get_frame_size
+    if _zstd_decompressor_class and _zstd_get_frame_size:
         return _zstd_decompressor_class
 
     global _importing_zstd
@@ -602,7 +603,10 @@ def _get_zstd_decompressor_class():
 
     _importing_zstd = True
     try:
-        from _zstd import ZstdDecompressor as _zstd_decompressor_class
+        from _zstd import (
+            ZstdDecompressor as _zstd_decompressor_class,
+            get_frame_size as _zstd_get_frame_size,
+        )
     except Exception:
         _bootstrap._verbose_message("zipimport: zstd UNAVAILABLE")
         raise ZipImportError("can't decompress data; zstd not available")
@@ -616,16 +620,24 @@ def _get_zstd_decompressor_class():
 def _zstd_decompress(data):
     # A simple version of compression.zstd.decompress() as we cannot import
     # that here as the stdlib itself could be being zipimported.
+    decomp_class = _get_zstd_decompressor_class()
+    data = memoryview(data)
+    if not data:
+        raise ZipImportError("zipimport: zstd compressed data ended before "
+                             "the end-of-stream marker")
     results = []
-    while True:
-        decomp = _get_zstd_decompressor_class()()
-        results.append(decomp.decompress(data))
+    while data:
+        try:
+            frame_size = _zstd_get_frame_size(data)
+        except Exception:
+            raise ZipImportError("zipimport: zstd compressed data ended before "
+                                 "the end-of-stream marker") from None
+        decomp = decomp_class()
+        results.append(decomp.decompress(data[:frame_size]))
         if not decomp.eof:
             raise ZipImportError("zipimport: zstd compressed data ended before "
                                  "the end-of-stream marker")
-        data = decomp.unused_data
-        if not data:
-            break
+        data = data[frame_size:]
     return b"".join(results)
 
 

Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst

@@ -0,0 +1 @@
+Improve performance of some :mod:`zipimport` archives that use Zstandard compression.