From 8f598e74650e711ef9429c69518c1b3b0aeaffc9 Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 12:12:29 +0200 Subject: [PATCH 1/7] test: add failing test for multi-frame zstd zip entries --- tests/util/test_zipfile_multiframe.py | 103 ++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 tests/util/test_zipfile_multiframe.py diff --git a/tests/util/test_zipfile_multiframe.py b/tests/util/test_zipfile_multiframe.py new file mode 100644 index 0000000000..5ce7ec17b6 --- /dev/null +++ b/tests/util/test_zipfile_multiframe.py @@ -0,0 +1,103 @@ +"""Tests for multi-frame zstd compression of ZIP entries. + +Producer-side fix: every ZIP_ZSTANDARD entry should be a multi-frame zstd +stream, capped at 200 MiB of input per frame, so JS decoders with per-frame +size limits (fzstd @ 256 MiB) can decode large inspect_ai .eval files. +""" + +from __future__ import annotations + +import zipfile +from pathlib import Path + +import pytest +import zstandard + +# Importing this module installs the zstd compression patches (both the +# zipfile_zstd delegation on Python < 3.14 and our multi-frame wrapper). +import inspect_ai._util.zipfile # noqa: F401 + +MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 # must match the value in _util/zipfile.py +ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" + + +def _moderately_compressible_payload(total_bytes: int) -> bytes: + """Build a deterministic payload that compresses at a realistic ratio. + + Pure random data bypasses zstd's match-finder (compressor emits raw + blocks); pure-constant data collapses to tiny output. We want something + in between so the compressed output is a non-trivial fraction of input — + close to the JSON-like eval workload we care about. + """ + fragments = [ + b'{"role": "assistant", "content": "The quick brown fox jumps over the lazy dog."}\n', + b'{"role": "user", "content": "What is the capital of France? Please answer in one sentence."}\n', + b'{"tool_calls": [{"name": "search", "arguments": {"query": "weather in paris"}}]}\n', + b'{"metadata": {"timestamp": "2026-04-23T10:00:00Z", "model": "claude-opus-4-7"}}\n', + ] + chunk = b"".join(fragments) + n = (total_bytes // len(chunk)) + 1 + return (chunk * n)[:total_bytes] + + +@pytest.fixture(scope="module") +def large_payload() -> bytes: + """~450 MiB of moderately compressible JSON-like bytes.""" + return _moderately_compressible_payload(450 * 1024 * 1024) + + +def _read_raw_compressed_entry(zip_path: Path, entry_name: str) -> bytes: + """Return the raw compressed bytes stored for a zip entry (no decompression).""" + with zipfile.ZipFile(zip_path) as zf: + info = zf.getinfo(entry_name) + with open(zip_path, "rb") as f: + f.seek(info.header_offset) + # Skip local file header to reach the data. Using ZipFile's private + # helper is fragile; instead, decode the local header manually. + import struct + + sig_version_etc = f.read(30) + assert sig_version_etc[:4] == b"PK\x03\x04", "not a local file header" + name_len = struct.unpack(" None: + """A ~450 MiB payload should produce ≥3 zstd frames, each ≤ 200 MiB uncompressed.""" + zip_path = tmp_path / "large.zip" + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + zf.writestr("big.json", large_payload) + + raw = _read_raw_compressed_entry(zip_path, "big.json") + + magic_count = raw.count(ZSTD_MAGIC) + assert magic_count >= 3, ( + f"expected ≥3 zstd frames in a 450 MiB entry, found {magic_count}" + ) + + frame_sizes = [len(out) for _, out in _iter_frames(raw)] + assert len(frame_sizes) == magic_count, ( + f"magic count {magic_count} != frames walked {len(frame_sizes)}" + ) + for i, size in enumerate(frame_sizes): + assert size <= MAX_INPUT_PER_FRAME, ( + f"frame {i} decompressed to {size} bytes, exceeds cap {MAX_INPUT_PER_FRAME}" + ) From e7541b2141fd0d864f77f8a4b3472d42a36da848 Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 12:15:41 +0200 Subject: [PATCH 2/7] feat: cap zstd frames at 200 MiB input for JS-decoder compatibility Co-Authored-By: Claude Sonnet 4.6 --- src/inspect_ai/_util/zipfile.py | 81 +++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index 6215059d72..b4149b3c20 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -1,6 +1,23 @@ +"""ZIP file helpers and monkey-patches for zstd support. + +On Python < 3.14 we import ``zipfile_zstd`` to monkey-patch zstandard +compression into the stdlib ``zipfile`` module. Python 3.14+ handles zstd +natively via stdlib. + +Additionally, we install a second monkey-patch on ``zipfile._get_compressor`` +that caps each emitted zstd frame at ``_MAX_INPUT_PER_FRAME`` bytes of input. +Large single zstd frames (>= 256 MiB compressed) trigger an overflow bug in +the pure-JS ``fzstd`` decoder used by our viewers; multi-framing keeps every +frame well under that threshold. +""" + +from __future__ import annotations + +import functools import logging import sys import zipfile +from collections.abc import Callable from typing import Any logger = logging.getLogger(__name__) @@ -14,4 +31,68 @@ "compresslevel": None, } + +# 200 MiB. Well under fzstd's 256 MiB (2^28) compressed-frame overflow +# threshold, applied to *input* bytes which compress smaller. +_MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 + + +class _MultiFrameZstdCompressObj: + """A zstd compressobj that chunks its output into multiple frames. + + Wraps a ``zstandard`` compressobj and flushes it (finalizing the current + frame) and replaces it (starting a new frame) every + ``_MAX_INPUT_PER_FRAME`` bytes of input. Multi-frame zstd streams are + valid per spec -- any compliant decoder reads them transparently. + """ + + def __init__(self, factory: Callable[[], Any]) -> None: + self._factory = factory + self._obj = factory() + self._input_bytes = 0 + + def compress(self, data: bytes) -> bytes: + out = b"" + offset = 0 + while offset < len(data): + remaining_cap = _MAX_INPUT_PER_FRAME - self._input_bytes + chunk = data[offset : offset + remaining_cap] + out += self._obj.compress(chunk) + self._input_bytes += len(chunk) + offset += len(chunk) + if self._input_bytes >= _MAX_INPUT_PER_FRAME: + out += self._obj.flush() + self._obj = self._factory() + self._input_bytes = 0 + return out + + def flush(self) -> bytes: + return self._obj.flush() + + +def _install_multiframe_compressor() -> None: + """Install the multi-frame zstd wrapper on ``zipfile._get_compressor``. + + Idempotent. Delegates to whatever ``_get_compressor`` was installed before + us (stdlib on Py >= 3.14; ``zipfile_zstd``'s patched version on Py < 3.14), + so compression level and thread count are preserved. + """ + if getattr(zipfile, "_inspect_ai_multiframe_installed", False): + return + + original = zipfile._get_compressor # type: ignore[attr-defined] + + def patched(compress_type: int, compresslevel: int | None = None) -> Any: + if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] + factory = functools.partial(original, compress_type, compresslevel) + return _MultiFrameZstdCompressObj(factory) + return original(compress_type, compresslevel) + + zipfile._get_compressor = patched # type: ignore[attr-defined] + zipfile._inspect_ai_multiframe_installed = True # type: ignore[attr-defined] + + +_install_multiframe_compressor() + + __all__ = ["zipfile_compress_kwargs"] From 04164bd9eb610b6f234deccb9660bbb3be127dce Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 12:17:15 +0200 Subject: [PATCH 3/7] test: verify small zstd entries stay single-frame --- tests/util/test_zipfile_multiframe.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/util/test_zipfile_multiframe.py b/tests/util/test_zipfile_multiframe.py index 5ce7ec17b6..5d19c03a17 100644 --- a/tests/util/test_zipfile_multiframe.py +++ b/tests/util/test_zipfile_multiframe.py @@ -101,3 +101,16 @@ def test_large_entry_emits_multiple_capped_frames( assert size <= MAX_INPUT_PER_FRAME, ( f"frame {i} decompressed to {size} bytes, exceeds cap {MAX_INPUT_PER_FRAME}" ) + + +def test_small_entry_emits_single_frame(tmp_path: Path) -> None: + """A 1 MiB entry should still be exactly one zstd frame.""" + zip_path = tmp_path / "small.zip" + payload = _moderately_compressible_payload(1 * 1024 * 1024) + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + zf.writestr("small.json", payload) + + raw = _read_raw_compressed_entry(zip_path, "small.json") + assert raw.count(ZSTD_MAGIC) == 1, ( + f"expected exactly 1 frame for small entry, got {raw.count(ZSTD_MAGIC)}" + ) From 1c22a7136c869f7aeab146089532bd0d1b167385 Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 12:20:11 +0200 Subject: [PATCH 4/7] test: verify multi-frame zstd round-trip preserves bytes The round-trip test revealed that the decompressor path also needed fixing: the previous ZstdDecompressObjWrapper (from zipfile_zstd) used a decompressobj that stopped after the first zstd frame, leaving subsequent frames unread and causing a CRC mismatch on read-back. Add _MultiFrameZstdDecompressObj, which chains fresh inner decompressobj instances across frame boundaries, and patch _get_decompressor alongside _get_compressor so both directions handle multi-frame streams correctly. --- src/inspect_ai/_util/zipfile.py | 83 +++++++++++++++++++++++---- tests/util/test_zipfile_multiframe.py | 14 +++++ 2 files changed, 86 insertions(+), 11 deletions(-) diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index b4149b3c20..f919b56a23 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -70,29 +70,90 @@ def flush(self) -> bytes: return self._obj.flush() -def _install_multiframe_compressor() -> None: - """Install the multi-frame zstd wrapper on ``zipfile._get_compressor``. +class _MultiFrameZstdDecompressObj: + """A zstd decompressobj that transparently spans multiple frames. - Idempotent. Delegates to whatever ``_get_compressor`` was installed before - us (stdlib on Py >= 3.14; ``zipfile_zstd``'s patched version on Py < 3.14), - so compression level and thread count are preserved. + ``zstandard.ZstdDecompressor().decompressobj()`` stops at the first frame + boundary and marks ``eof=True``. When the compressor splits large entries + into multiple frames (see ``_MultiFrameZstdCompressObj``), the reader must + recognise the frame boundary, start a fresh inner decompressobj, and + continue until all compressed bytes have been consumed. + + The stdlib ``zipfile._read1`` path for non-deflate compression does: + + data = self._decompressor.decompress(data) + self._eof = self._decompressor.eof or self._compress_left <= 0 + + We return ``eof=False`` always so that ``self._eof`` is driven purely by + ``self._compress_left <= 0`` (all compressed bytes fed). Meanwhile, + ``decompress`` buffers leftover bytes from a completed frame and feeds them + into the next inner decompressobj. + """ + + def __init__(self) -> None: + import zstandard # local import — already a hard dep via zipfile_zstd + + self._dctx = zstandard.ZstdDecompressor() + self._obj = self._dctx.decompressobj() + self._pending: bytes = b"" + + def decompress(self, data: bytes) -> bytes: + self._pending += data + out = b"" + while self._pending: + result = self._obj.decompress(self._pending) + out += result + if self._obj.eof: + # Inner frame complete; unused_data holds the start of the next. + self._pending = self._obj.unused_data + self._obj = self._dctx.decompressobj() + else: + # All pending input consumed; wait for more. + self._pending = b"" + break + return out + + def flush(self) -> bytes: + return b"" + + @property + def eof(self) -> bool: + # Always False: let compress_left drive the outer EOF check. + return False + + +def _install_multiframe_patches() -> None: + """Install multi-frame zstd compressor and decompressor patches. + + Idempotent. Wraps whatever ``_get_compressor`` / ``_get_decompressor`` were + installed before us (stdlib on Py >= 3.14; ``zipfile_zstd``'s versions on + Py < 3.14), so compression level and thread count are preserved. """ if getattr(zipfile, "_inspect_ai_multiframe_installed", False): return - original = zipfile._get_compressor # type: ignore[attr-defined] + original_compressor = zipfile._get_compressor # type: ignore[attr-defined] + original_decompressor = zipfile._get_decompressor # type: ignore[attr-defined] - def patched(compress_type: int, compresslevel: int | None = None) -> Any: + def patched_compressor(compress_type: int, compresslevel: int | None = None) -> Any: if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] - factory = functools.partial(original, compress_type, compresslevel) + factory = functools.partial( + original_compressor, compress_type, compresslevel + ) return _MultiFrameZstdCompressObj(factory) - return original(compress_type, compresslevel) + return original_compressor(compress_type, compresslevel) + + def patched_decompressor(compress_type: int) -> Any: + if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] + return _MultiFrameZstdDecompressObj() + return original_decompressor(compress_type) - zipfile._get_compressor = patched # type: ignore[attr-defined] + zipfile._get_compressor = patched_compressor # type: ignore[attr-defined] + zipfile._get_decompressor = patched_decompressor # type: ignore[attr-defined] zipfile._inspect_ai_multiframe_installed = True # type: ignore[attr-defined] -_install_multiframe_compressor() +_install_multiframe_patches() __all__ = ["zipfile_compress_kwargs"] diff --git a/tests/util/test_zipfile_multiframe.py b/tests/util/test_zipfile_multiframe.py index 5d19c03a17..238c0b5461 100644 --- a/tests/util/test_zipfile_multiframe.py +++ b/tests/util/test_zipfile_multiframe.py @@ -114,3 +114,17 @@ def test_small_entry_emits_single_frame(tmp_path: Path) -> None: assert raw.count(ZSTD_MAGIC) == 1, ( f"expected exactly 1 frame for small entry, got {raw.count(ZSTD_MAGIC)}" ) + + +def test_large_entry_round_trip(tmp_path: Path, large_payload: bytes) -> None: + """Writing and reading back a multi-frame entry must preserve bytes exactly.""" + zip_path = tmp_path / "rt.zip" + with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + zf.writestr("big.json", large_payload) + + with zipfile.ZipFile(zip_path) as zf: + got = zf.read("big.json") + + assert got == large_payload, ( + f"round-trip mismatch: input {len(large_payload)} bytes, got {len(got)} bytes" + ) From 40ce57b851b4fdb4fe0b8649ec7bdd17c051de25 Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 13:23:32 +0200 Subject: [PATCH 5/7] perf: avoid O(n^2) concat and repeated compressor construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two issues on multi-frame writes: 1. The wrapper's compress() accumulated output with += b"...", which is O(n^2) across the 1 GiB+ compressed output of a 1.38 GiB entry — gigabytes of needless copying. Use a list + b"".join, and a memoryview over the input to avoid slice copies. 2. The factory delegated to zipfile_zstd's _get_compressor, which constructs a fresh ZstdCompressor(threads=12) per compressobj call. Share a single compressor across all frames of an entry. Before: ~9.5 s for the 1.38 GiB entry (~3x baseline). After: ~3.9 s (+25% vs single-frame baseline, matching the inherent zstd multi-frame dictionary-reset overhead). --- src/inspect_ai/_util/zipfile.py | 34 ++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index f919b56a23..90a83babab 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -13,7 +13,6 @@ from __future__ import annotations -import functools import logging import sys import zipfile @@ -52,19 +51,22 @@ def __init__(self, factory: Callable[[], Any]) -> None: self._input_bytes = 0 def compress(self, data: bytes) -> bytes: - out = b"" + view = memoryview(data) + pieces: list[bytes] = [] offset = 0 - while offset < len(data): + n = len(view) + while offset < n: remaining_cap = _MAX_INPUT_PER_FRAME - self._input_bytes - chunk = data[offset : offset + remaining_cap] - out += self._obj.compress(chunk) - self._input_bytes += len(chunk) - offset += len(chunk) + end = min(offset + remaining_cap, n) + chunk = view[offset:end] + pieces.append(self._obj.compress(chunk)) + self._input_bytes += end - offset + offset = end if self._input_bytes >= _MAX_INPUT_PER_FRAME: - out += self._obj.flush() + pieces.append(self._obj.flush()) self._obj = self._factory() self._input_bytes = 0 - return out + return b"".join(pieces) def flush(self) -> bytes: return self._obj.flush() @@ -137,10 +139,16 @@ def _install_multiframe_patches() -> None: def patched_compressor(compress_type: int, compresslevel: int | None = None) -> Any: if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] - factory = functools.partial( - original_compressor, compress_type, compresslevel - ) - return _MultiFrameZstdCompressObj(factory) + # Share one ``ZstdCompressor`` across all frames of the entry. + # Delegating to ``original_compressor`` would instead create a + # fresh ``ZstdCompressor(threads=12)`` per frame, re-initialising + # its thread pool every 200 MiB. Matches zipfile_zstd's defaults + # (level=3, threads=12). + import zstandard + + level = 3 if compresslevel is None else compresslevel + compressor = zstandard.ZstdCompressor(level=level, threads=12) + return _MultiFrameZstdCompressObj(compressor.compressobj) return original_compressor(compress_type, compresslevel) def patched_decompressor(compress_type: int) -> Any: From 496bd1beb12f5579aeb1fd67a3e7d883286f82ec Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 14:09:40 +0200 Subject: [PATCH 6/7] refactor: clean up zstd patching - Resolve zipfile.ZIP_ZSTANDARD once at import time into _ZIP_ZSTANDARD, collapsing three # type: ignore suppressions into one and surfacing a missing attribute at import time instead of at call time. - Extract the threads count zipfile_zstd hardcodes (12) into _ZSTD_THREADS with a comment pointing at the source of truth, so our value and theirs can't silently drift apart. - Skip the trailing frame flush when the last compress() call landed exactly on a frame boundary; otherwise we appended an empty 9-byte zstd frame to such entries. --- src/inspect_ai/_util/zipfile.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index 90a83babab..3580117445 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -25,8 +25,12 @@ if sys.version_info < (3, 14): import zipfile_zstd # type: ignore[import-not-found, import-untyped] # noqa: F401 +# Resolve once so the rest of the module can use a plain int; also fails loudly +# at import time rather than at first call if the attribute is ever missing. +_ZIP_ZSTANDARD: int = zipfile.ZIP_ZSTANDARD # type: ignore[attr-defined] + zipfile_compress_kwargs: dict[str, Any] = { - "compression": zipfile.ZIP_ZSTANDARD, # type: ignore[attr-defined] + "compression": _ZIP_ZSTANDARD, "compresslevel": None, } @@ -35,6 +39,11 @@ # threshold, applied to *input* bytes which compress smaller. _MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 +# Matches zipfile_zstd's hardcoded default so multi-frame and single-frame zip +# writes use the same thread count. If zipfile_zstd ever changes its default, +# update here too. +_ZSTD_THREADS = 12 + class _MultiFrameZstdCompressObj: """A zstd compressobj that chunks its output into multiple frames. @@ -69,6 +78,11 @@ def compress(self, data: bytes) -> bytes: return b"".join(pieces) def flush(self) -> bytes: + # If the last ``compress()`` call landed exactly on the frame boundary, + # ``self._obj`` was replaced with a fresh compressobj that has received + # no bytes. Flushing it would append an empty 9-byte trailing frame. + if self._input_bytes == 0: + return b"" return self._obj.flush() @@ -138,21 +152,20 @@ def _install_multiframe_patches() -> None: original_decompressor = zipfile._get_decompressor # type: ignore[attr-defined] def patched_compressor(compress_type: int, compresslevel: int | None = None) -> Any: - if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] + if compress_type == _ZIP_ZSTANDARD: # Share one ``ZstdCompressor`` across all frames of the entry. # Delegating to ``original_compressor`` would instead create a - # fresh ``ZstdCompressor(threads=12)`` per frame, re-initialising - # its thread pool every 200 MiB. Matches zipfile_zstd's defaults - # (level=3, threads=12). + # fresh ``ZstdCompressor(threads=N)`` per frame, re-initialising + # its thread pool every 200 MiB. import zstandard level = 3 if compresslevel is None else compresslevel - compressor = zstandard.ZstdCompressor(level=level, threads=12) + compressor = zstandard.ZstdCompressor(level=level, threads=_ZSTD_THREADS) return _MultiFrameZstdCompressObj(compressor.compressobj) return original_compressor(compress_type, compresslevel) def patched_decompressor(compress_type: int) -> Any: - if compress_type == zipfile.ZIP_ZSTANDARD: # type: ignore[attr-defined] + if compress_type == _ZIP_ZSTANDARD: return _MultiFrameZstdDecompressObj() return original_decompressor(compress_type) From 883494b2b3ef32dbab9a561c59ab0d76dc1f1f8c Mon Sep 17 00:00:00 2001 From: Rasmus Faber-Espensen Date: Thu, 23 Apr 2026 14:42:42 +0200 Subject: [PATCH 7/7] fix: satisfy mypy on Python 3.10 / 3.11 - Type the compressor and decompressor factories with the concrete zstandard classes instead of Any, so self._obj.flush() no longer leaks Any into a function declared to return bytes. - In the test file, resolve ZIP_ZSTANDARD via getattr like the existing test_async_zip.py does, since on Python < 3.14 stdlib zipfile does not declare the attribute and type checkers don't see zipfile_zstd's runtime monkey-patch. --- src/inspect_ai/_util/zipfile.py | 13 ++++++++----- tests/util/test_zipfile_multiframe.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/inspect_ai/_util/zipfile.py b/src/inspect_ai/_util/zipfile.py index 3580117445..fe3dc51a21 100644 --- a/src/inspect_ai/_util/zipfile.py +++ b/src/inspect_ai/_util/zipfile.py @@ -17,7 +17,10 @@ import sys import zipfile from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import zstandard logger = logging.getLogger(__name__) @@ -54,7 +57,7 @@ class _MultiFrameZstdCompressObj: valid per spec -- any compliant decoder reads them transparently. """ - def __init__(self, factory: Callable[[], Any]) -> None: + def __init__(self, factory: Callable[[], zstandard.ZstdCompressionObj]) -> None: self._factory = factory self._obj = factory() self._input_bytes = 0 @@ -107,10 +110,10 @@ class _MultiFrameZstdDecompressObj: """ def __init__(self) -> None: - import zstandard # local import — already a hard dep via zipfile_zstd + import zstandard as zstd # local import -- already a hard dep - self._dctx = zstandard.ZstdDecompressor() - self._obj = self._dctx.decompressobj() + self._dctx: zstandard.ZstdDecompressor = zstd.ZstdDecompressor() + self._obj: zstandard.ZstdDecompressionObj = self._dctx.decompressobj() self._pending: bytes = b"" def decompress(self, data: bytes) -> bytes: diff --git a/tests/util/test_zipfile_multiframe.py b/tests/util/test_zipfile_multiframe.py index 238c0b5461..70cbcfe157 100644 --- a/tests/util/test_zipfile_multiframe.py +++ b/tests/util/test_zipfile_multiframe.py @@ -20,6 +20,11 @@ MAX_INPUT_PER_FRAME = 200 * 1024 * 1024 # must match the value in _util/zipfile.py ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" +# zipfile_zstd adds ZIP_ZSTANDARD (93) to zipfile at import time; use getattr +# so type checkers on Python < 3.14 (where stdlib zipfile doesn't declare it) +# stay happy. +ZIP_ZSTANDARD: int = getattr(zipfile, "ZIP_ZSTANDARD", 93) + def _moderately_compressible_payload(total_bytes: int) -> bytes: """Build a deterministic payload that compresses at a realistic ratio. @@ -83,7 +88,7 @@ def test_large_entry_emits_multiple_capped_frames( ) -> None: """A ~450 MiB payload should produce ≥3 zstd frames, each ≤ 200 MiB uncompressed.""" zip_path = tmp_path / "large.zip" - with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: zf.writestr("big.json", large_payload) raw = _read_raw_compressed_entry(zip_path, "big.json") @@ -107,7 +112,7 @@ def test_small_entry_emits_single_frame(tmp_path: Path) -> None: """A 1 MiB entry should still be exactly one zstd frame.""" zip_path = tmp_path / "small.zip" payload = _moderately_compressible_payload(1 * 1024 * 1024) - with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: zf.writestr("small.json", payload) raw = _read_raw_compressed_entry(zip_path, "small.json") @@ -119,7 +124,7 @@ def test_small_entry_emits_single_frame(tmp_path: Path) -> None: def test_large_entry_round_trip(tmp_path: Path, large_payload: bytes) -> None: """Writing and reading back a multi-frame entry must preserve bytes exactly.""" zip_path = tmp_path / "rt.zip" - with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_ZSTANDARD) as zf: + with zipfile.ZipFile(zip_path, "w", compression=ZIP_ZSTANDARD) as zf: zf.writestr("big.json", large_payload) with zipfile.ZipFile(zip_path) as zf: