fix: accelerate archives with pigz when available (#14)

Use pigz for archive operations when available, stage directory unpack on the target filesystem, and reuse unpacked Mathlib .lake caches for offline setup.

Author: LooKeng

leanup/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Lean-zh Community"""
 __email__ = 'leanprover@outlook.com'
-__version__ = '0.3.0'
+__version__ = '0.3.1'
 
 from .repo import (
     RepoManager,

leanup/cli/cache_ops.py

@@ -8,7 +8,7 @@
 import requests
 
 from leanup.const import LEANUP_CACHE_DIR
-from leanup.ops.environment import safe_extract, tar_directory
+from leanup.ops.environment import extract_tar_gz, tar_directory
 from leanup.paths import cache_dir as leanup_cache_dir
 from leanup.repo.cache_server import run_cache_server
 from leanup.repo.mathlib_cache import MathlibCacheManager, normalize_lean_version, remove_path
@@ -31,9 +31,9 @@ def _extract_lake_archive(archive: Path, target_lake: Path) -> Path:
         raise ValueError(f"Archive not found: {archive}")
     parent = target_lake.parent
     parent.mkdir(parents=True, exist_ok=True)
-    with tempfile.TemporaryDirectory(prefix="leanup-lake-unpack-") as work:
+    with tempfile.TemporaryDirectory(prefix=".leanup-lake-unpack.", dir=parent) as work:
         temp_root = Path(work)
-        safe_extract(archive, temp_root)
+        extract_tar_gz(archive, temp_root)
         extracted = temp_root / ".lake"
         if not extracted.exists() or not extracted.is_dir():
             raise ValueError(f"Archive does not contain top-level .lake/ directory: {archive}")

leanup/ops/environment.py

@@ -2,6 +2,7 @@
 
 from pathlib import Path
 import os
+import shutil
 import subprocess
 import tarfile
 import tempfile
@@ -13,6 +14,27 @@
 from leanup.repo.mathlib_cache import normalize_lean_version, remove_path
 
 
+def has_parallel_gzip() -> bool:
+    return shutil.which("pigz") is not None and shutil.which("tar") is not None
+
+
+def validate_archive_paths(archive: Path, target_dir: Path) -> None:
+    target_dir = target_dir.resolve()
+    with tarfile.open(archive, "r:gz") as tar:
+        for member in tar.getmembers():
+            member_path = (target_dir / member.name).resolve()
+            if not str(member_path).startswith(str(target_dir)):
+                raise ValueError(f"Archive contains unsafe path: {member.name}")
+
+
+def extract_tar_gz(archive: Path, target_dir: Path) -> None:
+    validate_archive_paths(archive, target_dir)
+    if has_parallel_gzip():
+        subprocess.run(["tar", "-I", "pigz", "-xf", str(archive), "-C", str(target_dir)], check=True)
+        return
+    safe_extract(archive, target_dir)
+
+
 def download_to(url: str, output_file: Path) -> Path:
     output_file.parent.mkdir(parents=True, exist_ok=True)
     with tempfile.NamedTemporaryFile(dir=output_file.parent, prefix=f".{output_file.name}.", suffix=".tmp", delete=False) as handle:
@@ -61,14 +83,25 @@ def tar_directory(source_dir: Path, arcname: str, output_file: Path, exclude: se
     with tempfile.NamedTemporaryFile(dir=output_file.parent, prefix=f".{output_file.name}.", suffix=".tmp", delete=False) as handle:
         temp_output = Path(handle.name)
     try:
-        with tarfile.open(temp_output, "w:gz", dereference=False) as tar:
-            if exclude:
-                for child in sorted(source_dir.iterdir()):
-                    if child.name in exclude:
-                        continue
-                    tar.add(child, arcname=f"{arcname}/{child.name}", recursive=True)
-            else:
-                tar.add(source_dir, arcname=arcname, recursive=True)
+        if exclude or not has_parallel_gzip():
+            with tarfile.open(temp_output, "w:gz", dereference=False) as tar:
+                if exclude:
+                    for child in sorted(source_dir.iterdir()):
+                        if child.name in exclude:
+                            continue
+                        tar.add(child, arcname=f"{arcname}/{child.name}", recursive=True)
+                else:
+                    tar.add(source_dir, arcname=arcname, recursive=True)
+        else:
+            subprocess.run(
+                ["tar", "-I", "pigz", "-cf", str(temp_output), "-C", str(source_dir.parent), "--", source_dir.name],
+                check=True,
+            )
+            # Preserve the requested archive root name. External tar is only used when arcname equals source name or .lake.
+            if arcname != source_dir.name:
+                remove_path(temp_output)
+                with tarfile.open(temp_output, "w:gz", dereference=False) as tar:
+                    tar.add(source_dir, arcname=arcname, recursive=True)
         temp_output.replace(output_file)
         return output_file
     except Exception:
@@ -117,9 +150,10 @@ def get_elan(server: str | None = None) -> Path:
 def unpack_elan(archive: Path | None = None, target_home: Path | None = None) -> Path:
     archive_path = archive or elan_archive_path()
     target = target_home or elan_home()
-    with tempfile.TemporaryDirectory(prefix="leanup-elan-unpack-") as work:
+    target.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix=".leanup-elan-unpack.", dir=target.parent) as work:
         work_root = Path(work)
-        safe_extract(archive_path, work_root)
+        extract_tar_gz(archive_path, work_root)
         extracted = work_root / ".elan"
         if not extracted.exists():
             raise ValueError(f"Archive does not contain .elan/: {archive_path}")
@@ -160,9 +194,10 @@ def get_lean(version: str, server: str | None = None) -> Path:
 def unpack_lean(version: str, archive: Path | None = None, target_home: Path | None = None) -> Path:
     archive_path = archive or lean_archive_path(version)
     home = target_home or elan_home()
-    with tempfile.TemporaryDirectory(prefix="leanup-lean-unpack-") as work:
+    (home / "toolchains").mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix=".leanup-lean-unpack.", dir=home / "toolchains") as work:
         work_root = Path(work)
-        safe_extract(archive_path, work_root)
+        extract_tar_gz(archive_path, work_root)
         toolchains_root = work_root / ".elan" / "toolchains"
         candidates = [path for path in toolchains_root.iterdir() if path.is_dir()] if toolchains_root.exists() else []
         if len(candidates) != 1:

leanup/repo/project_setup.py

@@ -12,6 +12,7 @@
 from leanup.repo.mathlib_cache import MathlibCacheManager, normalize_lean_version, remove_path
 from leanup.repo.manager import LeanRepo
 from leanup.utils.basic import working_directory
+from leanup.paths import cache_dir as leanup_cache_dir
 from leanup.utils.custom_logger import setup_logger
 
 logger = setup_logger("project_setup")
@@ -108,10 +109,13 @@ def setup(self, config: SetupConfig) -> SetupResult:
             cache_dir = config.mathlib_cache_dir if config.mathlib else None
 
             if config.mathlib and config.resolved_dependency_mode in {"symlink", "copy"}:
-                logger.info("Checking reusable mathlib package cache")
-                used_cache = self._prepare_mathlib_cache(config, project_dir)
-                if used_cache:
-                    self._write_manifest_from_packages(config, project_dir)
+                logger.info("Checking reusable local .lake cache")
+                used_cache = self._prepare_local_lake_cache(config, project_dir)
+                if not used_cache:
+                    logger.info("Checking reusable mathlib package cache")
+                    used_cache = self._prepare_mathlib_cache(config, project_dir)
+                    if used_cache:
+                        self._write_manifest_from_packages(config, project_dir)
 
             if config.mathlib and self._should_run_lake_update(config, project_dir):
                 logger.info("Running lake update")
@@ -365,6 +369,24 @@ def _verify_mathlib_project(self, project_dir: Path) -> None:
             if probe.exists():
                 probe.unlink()
 
+    def _local_lake_cache_dir(self, lean_version: str) -> Path:
+        return leanup_cache_dir() / "local" / "mathlib" / normalize_lean_version(lean_version) / ".lake"
+
+    def _prepare_local_lake_cache(self, config: SetupConfig, project_dir: Path) -> bool:
+        lake_cache = self._local_lake_cache_dir(config.lean_version)
+        if not lake_cache.exists():
+            return False
+
+        project_lake = project_dir / ".lake"
+        remove_path(project_lake)
+        project_lake.parent.mkdir(parents=True, exist_ok=True)
+        if config.resolved_dependency_mode == "symlink":
+            project_lake.symlink_to(lake_cache, target_is_directory=True)
+        else:
+            shutil.copytree(lake_cache, project_lake, symlinks=True)
+        self._write_manifest_from_packages(config, project_dir)
+        return True
+
     def _prepare_mathlib_cache(self, config: SetupConfig, project_dir: Path) -> bool:
         cache_dir = self.cache_manager.ensure_local_cache(config.lean_version)
         if not cache_dir:

leanup/repo/toolchain_cache.py

@@ -12,6 +12,7 @@
 from leanup.const import LEANUP_CACHE_DIR
 from leanup.repo.elan import ElanManager
 from leanup.repo.mathlib_cache import normalize_lean_version, remove_path
+from leanup.ops.environment import extract_tar_gz
 from leanup.utils.custom_logger import setup_logger
 
 logger = setup_logger("toolchain_cache")
@@ -101,8 +102,7 @@ def unpack_base_archive(self, archive_path: Optional[Path] = None) -> Path:
         archive_path = archive_path or self.get_base_archive_path()
         with tempfile.TemporaryDirectory(prefix=".elan-base.", dir=self.elan_home.parent) as work:
             temp_root = Path(work)
-            with tarfile.open(archive_path, "r:gz") as tar:
-                self._safe_extract(tar, temp_root)
+            extract_tar_gz(archive_path, temp_root)
             extracted = temp_root / ".elan"
             if not extracted.exists():
                 raise ValueError(f"Archive does not contain top-level .elan/ directory: {archive_path}")
@@ -136,8 +136,7 @@ def unpack_toolchain_archive(self, version: str, archive_path: Optional[Path] =
         archive_path = archive_path or self.get_toolchain_archive_path(version)
         with tempfile.TemporaryDirectory(prefix=".elan-toolchain.", dir=self.elan_home.parent) as work:
             temp_root = Path(work)
-            with tarfile.open(archive_path, "r:gz") as tar:
-                self._safe_extract(tar, temp_root)
+            extract_tar_gz(archive_path, temp_root)
             toolchains_root = temp_root / ".elan" / "toolchains"
             toolchain_dirs = [path for path in toolchains_root.iterdir() if path.is_dir()] if toolchains_root.exists() else []
             if len(toolchain_dirs) != 1: