diff --git a/src/manage/install_command.py b/src/manage/install_command.py index bcfe9d8..89b589e 100644 --- a/src/manage/install_command.py +++ b/src/manage/install_command.py @@ -166,6 +166,7 @@ def validate_package(install, dest, *, delete=True): def extract_package(package, prefix, calculate_dest=Path, *, on_progress=None, repair=False): + import shutil import zipfile LOGGER.debug("Starting extract of %s to %s", package, prefix) @@ -205,8 +206,8 @@ def _calc(prefix, filename, calculate_dest=calculate_dest): warn_overwrite.append(dest) continue ensure_tree(dest) - with open(dest, "wb") as f: - f.write(zf.read(member)) + with zf.open(member) as source, open(dest, "wb") as f: + shutil.copyfileobj(source, f, length=10 * 1024 * 1024) on_progress(100) if warn_out_of_prefix: diff --git a/tests/test_install_command.py b/tests/test_install_command.py index 7e8991d..45b6cdd 100644 --- a/tests/test_install_command.py +++ b/tests/test_install_command.py @@ -10,6 +10,51 @@ from manage.logging import LOGGER +@pytest.mark.parametrize("suffix", [".zip", ".nupkg"]) +@pytest.mark.parametrize("repair", [False, True]) +def test_extract_package_streaming(tmp_path, monkeypatch, suffix, repair): + """Extract in bounded reads while preserving overwrite and repair behavior.""" + import zipfile + + package = tmp_path / ("package" + suffix) + prefix = tmp_path / "install" + prefix.mkdir() + existing = prefix / "existing.txt" + existing.write_bytes(b"original") + chunk_size = 10 * 1024 * 1024 + data = bytes(range(256)) * (chunk_size // 256 + 1) + archive_prefix = "tools/" if suffix == ".nupkg" else "" + with zipfile.ZipFile(package, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr(archive_prefix + "nested/data.bin", data) + zf.writestr(archive_prefix + "empty.txt", b"") + zf.writestr(archive_prefix + "existing.txt", b"replacement") + if suffix == ".nupkg": + zf.writestr("metadata.txt", b"ignored") + + reads = [] + original_read = zipfile.ZipExtFile.read + + def bounded_read(self, n=-1): + assert n == chunk_size + result = original_read(self, n) + if self.name == archive_prefix + "nested/data.bin": + reads.append(len(result)) + return result + + monkeypatch.setattr(zipfile.ZipExtFile, "read", bounded_read) + progress = [] + IC.extract_package(package, prefix, calculate_dest=Path, + on_progress=progress.append, repair=repair) + assert (prefix / "nested/data.bin").read_bytes() == data + assert (prefix / "empty.txt").read_bytes() == b"" + assert existing.read_bytes() == (b"replacement" if repair else b"original") + assert not (prefix / "metadata.txt").exists() + assert reads == [chunk_size, 256, 0] + assert progress[0] == 0 + assert 100 in progress + assert (None in progress) == (not repair) + + def test_print_cli_shortcuts(patched_installs, assert_log, monkeypatch, tmp_path): class Cmd: scratch = {}