feat: add option to set zip compression level

src/auditwheel/main_repair.py

@@ -2,6 +2,7 @@
 
 import argparse
 import logging
+import zlib
 from pathlib import Path
 
 from auditwheel.patcher import Patchelf
@@ -40,6 +41,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-compression-level",
+        action=EnvironmentDefault,
+        metavar="ZIP_COMPRESSION_LEVEL",
+        env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
+        dest="ZIP_COMPRESSION_LEVEL",
+        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,
@@ -197,6 +210,7 @@ def execute(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
             patcher=patcher,
             exclude=exclude,
             strip=args.STRIP,
+            zip_compression_level=args.ZIP_COMPRESSION_LEVEL,
         )
 
         if out_wheel is not None:

src/auditwheel/repair.py

@@ -41,7 +41,8 @@ def repair_wheel(
     update_tags: bool,
     patcher: ElfPatcher,
     exclude: frozenset[str],
-    strip: bool = False,
+    strip: bool,
+    zip_compression_level: int,
 ) -> Path | None:
     elf_data = get_wheel_elfdata(wheel_policy, wheel_path, exclude)
     external_refs_by_fn = elf_data.full_external_refs
@@ -57,6 +58,7 @@ def repair_wheel(
 
     with InWheelCtx(wheel_path) as ctx:
         ctx.out_wheel = out_dir / wheel_fname
+        ctx.zip_compression_level = zip_compression_level
 
         match = WHEEL_INFO_RE(wheel_fname)
         if not match:

src/auditwheel/tools.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import argparse
+import logging
 import os
 import subprocess
 import zipfile
@@ -11,6 +12,8 @@
 
 _T = TypeVar("_T")
 
+logger = logging.getLogger(__name__)
+
 
 def unique_by_index(sequence: Iterable[_T]) -> list[_T]:
     """unique elements in `sequence` in the order in which they occur
@@ -90,6 +93,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,9 +106,17 @@ 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:
+def dir2zip(
+    in_dir: Path,
+    zip_fname: Path,
+    zip_compression_level: int,
+    date_time: datetime | None,
+) -> None:
     """Make a zip file `zip_fname` with contents of directory `in_dir`
 
     The recorded filenames are relative to `in_dir`, so doing a standard zip
@@ -117,9 +129,14 @@ def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) ->
         Directory path containing files to go in the zip archive
     zip_fname : Path
         Filename of zip archive to write
+    zip_compression_level: int
+        zlib.Z_DEFAULT_COMPRESSION (-1 aka. level 6) balances speed and size.
+        zlib.Z_NO_COMPRESSION (O) for some test builds that needs no compression at all
+        zlib.Z_BEST_COMPRESSION (9) for bandwidth-constrained or large amount of downloads
     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 +157,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=zip_compression_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,15 +177,33 @@ 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:
+                try:
+                    self.env_default = type(self.env_default)
+                except Exception:
+                    self.option_strings = kwargs["option_strings"]
+                    args = {
+                        "value": self.env_default,
+                        "type": type,
+                        "env": self.env,
+                    }
+                    msg = (
+                        "invalid type: %(value)r from environment variable "
+                        "%(env)r cannot be converted to %(type)r"
+                    )
+                    raise argparse.ArgumentError(self, msg % args) from None
             default = self.env_default
-        if default:
-            required = False
-        if self.env_default and choices is not None and self.env_default not in choices:
+        if (
+            self.env_default is not None
+            and choices is not None
+            and self.env_default not in choices
+        ):
             self.option_strings = kwargs["option_strings"]
             args = {
                 "value": self.env_default,
@@ -178,7 +216,12 @@ def __init__(
             )
             raise argparse.ArgumentError(self, msg % args)
 
-        super().__init__(default=default, required=required, choices=choices, **kwargs)
+        if default is not None:
+            required = False
+
+        super().__init__(
+            default=default, required=required, choices=choices, type=type, **kwargs
+        )
 
     def __call__(
         self,

src/auditwheel/wheeltools.py

@@ -9,6 +9,7 @@
 import hashlib
 import logging
 import os
+import zlib
 from base64 import urlsafe_b64encode
 from collections.abc import Generator, Iterable
 from datetime import datetime, timezone
@@ -113,6 +114,7 @@ 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()
+        self.zip_compression_level = zlib.Z_DEFAULT_COMPRESSION
         super().__init__()
 
     def __enter__(self) -> Path:
@@ -131,7 +133,7 @@ def __exit__(
             timestamp = os.environ.get("SOURCE_DATE_EPOCH")
             if timestamp:
                 date_time = datetime.fromtimestamp(int(timestamp), tz=timezone.utc)
-            dir2zip(self.name, self.out_wheel, date_time)
+            dir2zip(self.name, self.out_wheel, self.zip_compression_level, date_time)
         return super().__exit__(exc, value, tb)
 
 

tests/integration/test_bundled_wheels.py

@@ -134,6 +134,7 @@ def test_wheel_source_date_epoch(tmp_path, monkeypatch):
         WHEEL_FILE=[wheel_path],
         EXCLUDE=[],
         DISABLE_ISA_EXT_CHECK=False,
+        ZIP_COMPRESSION_LEVEL=6,
         cmd="repair",
         func=Mock(),
         prog="auditwheel",

tests/unit/test_tools.py

@@ -3,6 +3,7 @@
 import argparse
 import lzma
 import zipfile
+import zlib
 from pathlib import Path
 
 import pytest
@@ -19,7 +20,7 @@
         ("manylinux2010", "linux", "linux"),
     ],
 )
-def test_environment_action(
+def test_plat_environment_action(
     monkeypatch: pytest.MonkeyPatch,
     environ: str | None,
     passed: str | None,
@@ -44,7 +45,50 @@ def test_environment_action(
     assert expected == args.PLAT
 
 
-def test_environment_action_invalid_env(monkeypatch: pytest.MonkeyPatch) -> None:
+_all_zip_level: list[int] = list(
+    range(zlib.Z_NO_COMPRESSION, zlib.Z_BEST_COMPRESSION + 1)
+)
+
+
+@pytest.mark.parametrize(
+    ("environ", "passed", "expected"),
+    [
+        (None, None, -1),
+        (0, None, 0),
+        (0, 1, 1),
+        (6, 1, 1),
+    ],
+)
+def test_zip_environment_action(
+    monkeypatch: pytest.MonkeyPatch,
+    environ: int | None,
+    passed: int | None,
+    expected: int,
+) -> None:
+    choices = _all_zip_level
+    argv = []
+    if passed is not None:
+        argv = ["--zip-compression-level", str(passed)]
+    if environ is not None:
+        monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", str(environ))
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "-z",
+        "--zip-compression-level",
+        action=EnvironmentDefault,
+        metavar="zip",
+        env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
+        dest="zip",
+        type=int,
+        help="Compress level to be used to create zip file.",
+        choices=choices,
+        default=zlib.Z_DEFAULT_COMPRESSION,
+    )
+    args = p.parse_args(argv)
+    assert expected == args.zip
+
+
+def test_environment_action_invalid_plat_env(monkeypatch: pytest.MonkeyPatch) -> None:
     choices = ["linux", "manylinux1", "manylinux2010"]
     monkeypatch.setenv("AUDITWHEEL_PLAT", "foo")
     p = argparse.ArgumentParser()
@@ -59,6 +103,39 @@ def test_environment_action_invalid_env(monkeypatch: pytest.MonkeyPatch) -> None
         )
 
 
+def test_environment_action_invalid_zip_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    choices = _all_zip_level
+    monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", "foo")
+    p = argparse.ArgumentParser()
+    with pytest.raises(argparse.ArgumentError):
+        p.add_argument(
+            "-z",
+            "--zip-compression-level",
+            action=EnvironmentDefault,
+            metavar="zip",
+            env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
+            dest="zip",
+            type=int,
+            help="Compress level to be used to create zip file.",
+            choices=choices,
+            default=zlib.Z_DEFAULT_COMPRESSION,
+        )
+    monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", "10")
+    with pytest.raises(argparse.ArgumentError):
+        p.add_argument(
+            "-z",
+            "--zip-compression-level",
+            action=EnvironmentDefault,
+            metavar="zip",
+            env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
+            dest="zip",
+            type=int,
+            help="Compress level to be used to create zip file.",
+            choices=choices,
+            default=zlib.Z_DEFAULT_COMPRESSION,
+        )
+
+
 def _write_test_permissions_zip(path: Path) -> None:
     source_zip_xz = Path(__file__).parent / "test-permissions.zip.xz"
     with lzma.open(source_zip_xz) as f:
@@ -92,7 +169,7 @@ def test_zip2dir_round_trip_permissions(tmp_path: Path) -> None:
     _write_test_permissions_zip(source_zip)
     extract_path = tmp_path / "unzip2"
     zip2dir(source_zip, tmp_path / "unzip1")
-    dir2zip(tmp_path / "unzip1", tmp_path / "tmp.zip")
+    dir2zip(tmp_path / "unzip1", tmp_path / "tmp.zip", zlib.Z_DEFAULT_COMPRESSION, None)
     zip2dir(tmp_path / "tmp.zip", extract_path)
     _check_permissions(extract_path)
 
@@ -104,7 +181,7 @@ def test_dir2zip_deflate(tmp_path: Path) -> None:
     input_file = input_dir / "zeros.bin"
     input_file.write_bytes(buffer)
     output_file = tmp_path / "ouput.zip"
-    dir2zip(input_dir, output_file)
+    dir2zip(input_dir, output_file, zlib.Z_DEFAULT_COMPRESSION, None)
     assert output_file.stat().st_size < len(buffer) / 4
 
 
@@ -117,7 +194,7 @@ def test_dir2zip_folders(tmp_path: Path) -> None:
     empty_folder = input_dir / "dummy" / "empty"
     empty_folder.mkdir(parents=True)
     output_file = tmp_path / "output.zip"
-    dir2zip(input_dir, output_file)
+    dir2zip(input_dir, output_file, zlib.Z_DEFAULT_COMPRESSION, None)
     expected_dirs = {"dummy/", "dummy/empty/", "dummy-1.0.dist-info/"}
     with zipfile.ZipFile(output_file, "r") as z:
         assert len(z.filelist) == 4