https://github.com/python/cpython/commit/1196c345377181de5510c5adece5d1ab7e010973 commit: 1196c345377181de5510c5adece5d1ab7e010973 branch: 3.15 author: Miss Islington (bot) <[email protected]> committer: emmatyping <[email protected]> date: 2026-10-24T14:55:06-07:00 summary:
[3.15] gh-155152: Improve performance of concatenated zstd frames in zipimport (GH-155154) (#160565) gh-155152: Improve performance of concatenated zstd frames in zipimport (GH-155154) (cherry picked from commit 95e088ed4964b438bc404c378f9f7915878ee5f7) Co-authored-by: Pengyu Lee <[email protected]> Co-authored-by: Emma Smith <[email protected]> files: A Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst M Lib/test/test_zipimport.py M Lib/zipimport.py diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py index b30068d696c7eb..94875d7153f9c5 100644 --- a/Lib/test/test_zipimport.py +++ b/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): diff --git a/Lib/zipimport.py b/Lib/zipimport.py index f066fd71961fa5..09a85f6d7f58e8 100644 --- a/Lib/zipimport.py +++ b/Lib/zipimport.py @@ -590,12 +590,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 @@ -607,7 +608,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") @@ -621,16 +625,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) diff --git a/Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst b/Misc/NEWS.d/next/Library/2026-08-04-16-38-59.gh-issue-155152.qN1EEN.rst new file mode 100644 index 00000000000000..c7cff1866cc51f --- /dev/null +++ b/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. _______________________________________________ Python-checkins mailing list -- [email protected] To unsubscribe send an email to [email protected] https://mail.python.org/mailman3//lists/python-checkins.python.org Member address: [email protected]
