diff --git a/src/auditwheel/main_repair.py b/src/auditwheel/main_repair.py index 65e356ee..47c09157 100644 --- a/src/auditwheel/main_repair.py +++ b/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: diff --git a/src/auditwheel/repair.py b/src/auditwheel/repair.py index 964676cc..82cf6084 100644 --- a/src/auditwheel/repair.py +++ b/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: diff --git a/src/auditwheel/tools.py b/src/auditwheel/tools.py index 0592ad2d..70d30ab0 100644 --- a/src/auditwheel/tools.py +++ b/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, diff --git a/src/auditwheel/wheeltools.py b/src/auditwheel/wheeltools.py index 828fbeee..3435a873 100644 --- a/src/auditwheel/wheeltools.py +++ b/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) diff --git a/tests/integration/test_bundled_wheels.py b/tests/integration/test_bundled_wheels.py index 33d76a3d..4605ec54 100644 --- a/tests/integration/test_bundled_wheels.py +++ b/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", diff --git a/tests/unit/test_tools.py b/tests/unit/test_tools.py index 64874af0..d5c022b9 100644 --- a/tests/unit/test_tools.py +++ b/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