Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/manage/install_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
45 changes: 45 additions & 0 deletions tests/test_install_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down