chore: Add argument and env to specific compress level

Author: oraluben

src/auditwheel/main_repair.py

@@ -2,10 +2,12 @@
 
 import argparse
 import logging
+import zlib
 from pathlib import Path
 
 from auditwheel.patcher import Patchelf
 
+from . import tools
 from .policy import WheelPolicies
 from .tools import EnvironmentDefault
 
@@ -40,6 +42,18 @@ def configure_parser(sub_parsers) -> None:  # type: ignore[no-untyped-def]
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
     parser.add_argument("WHEEL_FILE", type=Path, help="Path to wheel file.", nargs="+")
+    parser.add_argument(
+        "-z",
+        "--zip-level",
+        action=EnvironmentDefault,
+        metavar="zip",
+        env="AUDITWHEEL_ZIP_LEVEL",
+        dest="zip",
+        type=int,
+        help="Compress level to be used to create zip file.",
+        choices=list(range(zlib.Z_NO_COMPRESSION, zlib.Z_BEST_COMPRESSION + 1)),
+        default=zlib.Z_DEFAULT_COMPRESSION,
+    )
     parser.add_argument(
         "--plat",
         action=EnvironmentDefault,
@@ -119,6 +133,7 @@ def execute(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
     wheel_dir: Path = args.WHEEL_DIR.absolute()
     wheel_files: list[Path] = args.WHEEL_FILE
     wheel_policy = WheelPolicies()
+    tools._COMPRESS_LEVEL = args.zip
 
     for wheel_file in wheel_files:
         if not wheel_file.is_file():

src/auditwheel/tools.py

@@ -1,16 +1,27 @@
 from __future__ import annotations
 
 import argparse
+import logging
 import os
 import subprocess
 import zipfile
+import zlib
 from collections.abc import Generator, Iterable
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, TypeVar
 
 _T = TypeVar("_T")
 
+logger = logging.getLogger(__name__)
+
+# Default: zlib.Z_DEFAULT_COMPRESSION (-1 aka. level 6) balances speed and size.
+# Maintained for typical builds where iteration speed outweighs distribution savings.
+# Override via AUDITWHEEL_ZIP_LEVEL/--zip-level for:
+# - some test builds that needs no compression at all (0)
+# - bandwidth-constrained or large amount of downloads (9)
+_COMPRESS_LEVEL = zlib.Z_DEFAULT_COMPRESSION
+
 
 def unique_by_index(sequence: Iterable[_T]) -> list[_T]:
     """unique elements in `sequence` in the order in which they occur
@@ -90,6 +101,7 @@ def zip2dir(zip_fname: Path, out_dir: Path) -> None:
     out_dir : str
         Directory path containing files to go in the zip archive
     """
+    start = datetime.now()
     with zipfile.ZipFile(zip_fname, "r") as z:
         for name in z.namelist():
             member = z.getinfo(name)
@@ -102,6 +114,9 @@ def zip2dir(zip_fname: Path, out_dir: Path) -> None:
                 attr &= 511  # only keep permission bits
                 attr |= 6 << 6  # at least read/write for current user
                 os.chmod(extracted_path, attr)
+    logger.debug(
+        "zip2dir from %s to %s takes %s", zip_fname, out_dir, datetime.now() - start
+    )
 
 
 def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) -> None:
@@ -120,6 +135,7 @@ def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) ->
     date_time : Optional[datetime]
         Time stamp to set on each file in the archive
     """
+    start = datetime.now()
     in_dir = in_dir.resolve(strict=True)
     if date_time is None:
         st = in_dir.stat()
@@ -140,7 +156,10 @@ def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) ->
                 zinfo.date_time = date_time_args
                 zinfo.compress_type = compression
                 with open(fname, "rb") as fp:
-                    z.writestr(zinfo, fp.read())
+                    z.writestr(zinfo, fp.read(), compresslevel=_COMPRESS_LEVEL)
+    logger.debug(
+        "dir2zip from %s to %s takes %s", in_dir, zip_fname, datetime.now() - start
+    )
 
 
 def tarbz2todir(tarbz2_fname: Path, out_dir: Path) -> None:
@@ -157,11 +176,14 @@ def __init__(
         required: bool = True,
         default: str | None = None,
         choices: Iterable[str] | None = None,
+        type: type | None = None,
         **kwargs: Any,
     ) -> None:
         self.env_default = os.environ.get(env)
         self.env = env
         if self.env_default:
+            if type:
+                self.env_default = type(self.env_default)
             default = self.env_default
         if default:
             required = False

src/auditwheel/wheeltools.py

@@ -15,7 +15,9 @@
 from itertools import product
 from os.path import splitext
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from types import TracebackType
+from typing import Any, ClassVar
 
 from packaging.utils import parse_wheel_filename
 
@@ -98,8 +100,13 @@ class InWheel(InTemporaryDirectory):
     On entering, you'll find yourself in the root tree of the wheel.  If you've
     asked for an output wheel, then on exit we'll rewrite the wheel record and
     pack stuff up for you.
+
+    If `out_wheel` is None, we assume the wheel won't be modified and we can
+    cache the unpacked wheel for future use.
     """
 
+    _whl_cache: ClassVar[dict[Path, TemporaryDirectory[Any]]] = {}
+
     def __init__(self, in_wheel: Path, out_wheel: Path | None = None) -> None:
         """Initialize in-wheel context manager
 
@@ -113,9 +120,35 @@ def __init__(self, in_wheel: Path, out_wheel: Path | None = None) -> None:
         """
         self.in_wheel = in_wheel.absolute()
         self.out_wheel = None if out_wheel is None else out_wheel.absolute()
-        super().__init__()
+        self.read_only = out_wheel is None
+        self.use_cache = self.in_wheel in self._whl_cache
+        if self.use_cache and not Path(self._whl_cache[self.in_wheel].name).exists():
+            self.use_cache = False
+            logger.debug(
+                "Wheel ctx %s for %s is no longer valid",
+                self._whl_cache.pop(self.in_wheel),
+                self.in_wheel,
+            )
+
+        if self.use_cache:
+            logger.debug(
+                "Reuse %s for %s", self._whl_cache[self.in_wheel], self.in_wheel
+            )
+            self._tmpdir = self._whl_cache[self.in_wheel]
+            if not self.read_only:
+                self._whl_cache.pop(self.in_wheel)
+        else:
+            super().__init__()
+            if self.read_only:
+                self._whl_cache[self.in_wheel] = self._tmpdir
 
     def __enter__(self) -> Path:
+        if self.use_cache or self.read_only:
+            if not self.use_cache:
+                zip2dir(self.in_wheel, self.name)
+            self._pwd = Path.cwd()
+            os.chdir(self.name)
+            return Path(self.name)
         zip2dir(self.in_wheel, self.name)
         return super().__enter__()
 
@@ -132,6 +165,16 @@ def __exit__(
             if timestamp:
                 date_time = datetime.fromtimestamp(int(timestamp), tz=timezone.utc)
             dir2zip(self.name, self.out_wheel, date_time)
+        if self.use_cache or self.read_only:
+            logger.debug(
+                "Exiting reused %s for %s",
+                self._whl_cache[self.in_wheel],
+                self.in_wheel,
+            )
+            os.chdir(self._pwd)
+            if not self.read_only:
+                super().__exit__(exc, value, tb)
+            return None
         return super().__exit__(exc, value, tb)