Stream package extraction in 10MB chunks to limit peak memory usage (#410)

Fixes #409

10MB was chosen so that files in a standard Python distribution are all still copied in a single step. However, custom packages with larger files are now streamed.

Author: joaogabriel15

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:

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 = {}