Skip to content

Commit 0c1b26b

Browse files
authored
feat: add option to set zip compression level (#547)
Add --zip-compression-level and AUDITWHEEL_ZIP_COMPRESSION_LEVEL env to specify zip compression level.
1 parent 2874cd3 commit 0c1b26b

File tree

6 files changed

+152
-13
lines changed

6 files changed

+152
-13
lines changed

src/auditwheel/main_repair.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import argparse
44
import logging
5+
import zlib
56
from pathlib import Path
67

78
from auditwheel.patcher import Patchelf
@@ -40,6 +41,18 @@ def configure_parser(sub_parsers) -> None: # type: ignore[no-untyped-def]
4041
formatter_class=argparse.RawDescriptionHelpFormatter,
4142
)
4243
parser.add_argument("WHEEL_FILE", type=Path, help="Path to wheel file.", nargs="+")
44+
parser.add_argument(
45+
"-z",
46+
"--zip-compression-level",
47+
action=EnvironmentDefault,
48+
metavar="ZIP_COMPRESSION_LEVEL",
49+
env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
50+
dest="ZIP_COMPRESSION_LEVEL",
51+
type=int,
52+
help="Compress level to be used to create zip file.",
53+
choices=list(range(zlib.Z_NO_COMPRESSION, zlib.Z_BEST_COMPRESSION + 1)),
54+
default=zlib.Z_DEFAULT_COMPRESSION,
55+
)
4356
parser.add_argument(
4457
"--plat",
4558
action=EnvironmentDefault,
@@ -197,6 +210,7 @@ def execute(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
197210
patcher=patcher,
198211
exclude=exclude,
199212
strip=args.STRIP,
213+
zip_compression_level=args.ZIP_COMPRESSION_LEVEL,
200214
)
201215

202216
if out_wheel is not None:

src/auditwheel/repair.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ def repair_wheel(
4141
update_tags: bool,
4242
patcher: ElfPatcher,
4343
exclude: frozenset[str],
44-
strip: bool = False,
44+
strip: bool,
45+
zip_compression_level: int,
4546
) -> Path | None:
4647
elf_data = get_wheel_elfdata(wheel_policy, wheel_path, exclude)
4748
external_refs_by_fn = elf_data.full_external_refs
@@ -57,6 +58,7 @@ def repair_wheel(
5758

5859
with InWheelCtx(wheel_path) as ctx:
5960
ctx.out_wheel = out_dir / wheel_fname
61+
ctx.zip_compression_level = zip_compression_level
6062

6163
match = WHEEL_INFO_RE(wheel_fname)
6264
if not match:

src/auditwheel/tools.py

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import argparse
4+
import logging
45
import os
56
import subprocess
67
import zipfile
@@ -11,6 +12,8 @@
1112

1213
_T = TypeVar("_T")
1314

15+
logger = logging.getLogger(__name__)
16+
1417

1518
def unique_by_index(sequence: Iterable[_T]) -> list[_T]:
1619
"""unique elements in `sequence` in the order in which they occur
@@ -90,6 +93,7 @@ def zip2dir(zip_fname: Path, out_dir: Path) -> None:
9093
out_dir : str
9194
Directory path containing files to go in the zip archive
9295
"""
96+
start = datetime.now()
9397
with zipfile.ZipFile(zip_fname, "r") as z:
9498
for name in z.namelist():
9599
member = z.getinfo(name)
@@ -102,9 +106,17 @@ def zip2dir(zip_fname: Path, out_dir: Path) -> None:
102106
attr &= 511 # only keep permission bits
103107
attr |= 6 << 6 # at least read/write for current user
104108
os.chmod(extracted_path, attr)
109+
logger.debug(
110+
"zip2dir from %s to %s takes %s", zip_fname, out_dir, datetime.now() - start
111+
)
105112

106113

107-
def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) -> None:
114+
def dir2zip(
115+
in_dir: Path,
116+
zip_fname: Path,
117+
zip_compression_level: int,
118+
date_time: datetime | None,
119+
) -> None:
108120
"""Make a zip file `zip_fname` with contents of directory `in_dir`
109121
110122
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) ->
117129
Directory path containing files to go in the zip archive
118130
zip_fname : Path
119131
Filename of zip archive to write
132+
zip_compression_level: int
133+
zlib.Z_DEFAULT_COMPRESSION (-1 aka. level 6) balances speed and size.
134+
zlib.Z_NO_COMPRESSION (O) for some test builds that needs no compression at all
135+
zlib.Z_BEST_COMPRESSION (9) for bandwidth-constrained or large amount of downloads
120136
date_time : Optional[datetime]
121137
Time stamp to set on each file in the archive
122138
"""
139+
start = datetime.now()
123140
in_dir = in_dir.resolve(strict=True)
124141
if date_time is None:
125142
st = in_dir.stat()
@@ -140,7 +157,10 @@ def dir2zip(in_dir: Path, zip_fname: Path, date_time: datetime | None = None) ->
140157
zinfo.date_time = date_time_args
141158
zinfo.compress_type = compression
142159
with open(fname, "rb") as fp:
143-
z.writestr(zinfo, fp.read())
160+
z.writestr(zinfo, fp.read(), compresslevel=zip_compression_level)
161+
logger.debug(
162+
"dir2zip from %s to %s takes %s", in_dir, zip_fname, datetime.now() - start
163+
)
144164

145165

146166
def tarbz2todir(tarbz2_fname: Path, out_dir: Path) -> None:
@@ -157,15 +177,33 @@ def __init__(
157177
required: bool = True,
158178
default: str | None = None,
159179
choices: Iterable[str] | None = None,
180+
type: type | None = None,
160181
**kwargs: Any,
161182
) -> None:
162183
self.env_default = os.environ.get(env)
163184
self.env = env
164185
if self.env_default:
186+
if type:
187+
try:
188+
self.env_default = type(self.env_default)
189+
except Exception:
190+
self.option_strings = kwargs["option_strings"]
191+
args = {
192+
"value": self.env_default,
193+
"type": type,
194+
"env": self.env,
195+
}
196+
msg = (
197+
"invalid type: %(value)r from environment variable "
198+
"%(env)r cannot be converted to %(type)r"
199+
)
200+
raise argparse.ArgumentError(self, msg % args) from None
165201
default = self.env_default
166-
if default:
167-
required = False
168-
if self.env_default and choices is not None and self.env_default not in choices:
202+
if (
203+
self.env_default is not None
204+
and choices is not None
205+
and self.env_default not in choices
206+
):
169207
self.option_strings = kwargs["option_strings"]
170208
args = {
171209
"value": self.env_default,
@@ -178,7 +216,12 @@ def __init__(
178216
)
179217
raise argparse.ArgumentError(self, msg % args)
180218

181-
super().__init__(default=default, required=required, choices=choices, **kwargs)
219+
if default is not None:
220+
required = False
221+
222+
super().__init__(
223+
default=default, required=required, choices=choices, type=type, **kwargs
224+
)
182225

183226
def __call__(
184227
self,

src/auditwheel/wheeltools.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import hashlib
1010
import logging
1111
import os
12+
import zlib
1213
from base64 import urlsafe_b64encode
1314
from collections.abc import Generator, Iterable
1415
from datetime import datetime, timezone
@@ -113,6 +114,7 @@ def __init__(self, in_wheel: Path, out_wheel: Path | None = None) -> None:
113114
"""
114115
self.in_wheel = in_wheel.absolute()
115116
self.out_wheel = None if out_wheel is None else out_wheel.absolute()
117+
self.zip_compression_level = zlib.Z_DEFAULT_COMPRESSION
116118
super().__init__()
117119

118120
def __enter__(self) -> Path:
@@ -131,7 +133,7 @@ def __exit__(
131133
timestamp = os.environ.get("SOURCE_DATE_EPOCH")
132134
if timestamp:
133135
date_time = datetime.fromtimestamp(int(timestamp), tz=timezone.utc)
134-
dir2zip(self.name, self.out_wheel, date_time)
136+
dir2zip(self.name, self.out_wheel, self.zip_compression_level, date_time)
135137
return super().__exit__(exc, value, tb)
136138

137139

tests/integration/test_bundled_wheels.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ def test_wheel_source_date_epoch(tmp_path, monkeypatch):
134134
WHEEL_FILE=[wheel_path],
135135
EXCLUDE=[],
136136
DISABLE_ISA_EXT_CHECK=False,
137+
ZIP_COMPRESSION_LEVEL=6,
137138
cmd="repair",
138139
func=Mock(),
139140
prog="auditwheel",

tests/unit/test_tools.py

Lines changed: 82 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import argparse
44
import lzma
55
import zipfile
6+
import zlib
67
from pathlib import Path
78

89
import pytest
@@ -19,7 +20,7 @@
1920
("manylinux2010", "linux", "linux"),
2021
],
2122
)
22-
def test_environment_action(
23+
def test_plat_environment_action(
2324
monkeypatch: pytest.MonkeyPatch,
2425
environ: str | None,
2526
passed: str | None,
@@ -44,7 +45,50 @@ def test_environment_action(
4445
assert expected == args.PLAT
4546

4647

47-
def test_environment_action_invalid_env(monkeypatch: pytest.MonkeyPatch) -> None:
48+
_all_zip_level: list[int] = list(
49+
range(zlib.Z_NO_COMPRESSION, zlib.Z_BEST_COMPRESSION + 1)
50+
)
51+
52+
53+
@pytest.mark.parametrize(
54+
("environ", "passed", "expected"),
55+
[
56+
(None, None, -1),
57+
(0, None, 0),
58+
(0, 1, 1),
59+
(6, 1, 1),
60+
],
61+
)
62+
def test_zip_environment_action(
63+
monkeypatch: pytest.MonkeyPatch,
64+
environ: int | None,
65+
passed: int | None,
66+
expected: int,
67+
) -> None:
68+
choices = _all_zip_level
69+
argv = []
70+
if passed is not None:
71+
argv = ["--zip-compression-level", str(passed)]
72+
if environ is not None:
73+
monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", str(environ))
74+
p = argparse.ArgumentParser()
75+
p.add_argument(
76+
"-z",
77+
"--zip-compression-level",
78+
action=EnvironmentDefault,
79+
metavar="zip",
80+
env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
81+
dest="zip",
82+
type=int,
83+
help="Compress level to be used to create zip file.",
84+
choices=choices,
85+
default=zlib.Z_DEFAULT_COMPRESSION,
86+
)
87+
args = p.parse_args(argv)
88+
assert expected == args.zip
89+
90+
91+
def test_environment_action_invalid_plat_env(monkeypatch: pytest.MonkeyPatch) -> None:
4892
choices = ["linux", "manylinux1", "manylinux2010"]
4993
monkeypatch.setenv("AUDITWHEEL_PLAT", "foo")
5094
p = argparse.ArgumentParser()
@@ -59,6 +103,39 @@ def test_environment_action_invalid_env(monkeypatch: pytest.MonkeyPatch) -> None
59103
)
60104

61105

106+
def test_environment_action_invalid_zip_env(monkeypatch: pytest.MonkeyPatch) -> None:
107+
choices = _all_zip_level
108+
monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", "foo")
109+
p = argparse.ArgumentParser()
110+
with pytest.raises(argparse.ArgumentError):
111+
p.add_argument(
112+
"-z",
113+
"--zip-compression-level",
114+
action=EnvironmentDefault,
115+
metavar="zip",
116+
env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
117+
dest="zip",
118+
type=int,
119+
help="Compress level to be used to create zip file.",
120+
choices=choices,
121+
default=zlib.Z_DEFAULT_COMPRESSION,
122+
)
123+
monkeypatch.setenv("AUDITWHEEL_ZIP_COMPRESSION_LEVEL", "10")
124+
with pytest.raises(argparse.ArgumentError):
125+
p.add_argument(
126+
"-z",
127+
"--zip-compression-level",
128+
action=EnvironmentDefault,
129+
metavar="zip",
130+
env="AUDITWHEEL_ZIP_COMPRESSION_LEVEL",
131+
dest="zip",
132+
type=int,
133+
help="Compress level to be used to create zip file.",
134+
choices=choices,
135+
default=zlib.Z_DEFAULT_COMPRESSION,
136+
)
137+
138+
62139
def _write_test_permissions_zip(path: Path) -> None:
63140
source_zip_xz = Path(__file__).parent / "test-permissions.zip.xz"
64141
with lzma.open(source_zip_xz) as f:
@@ -92,7 +169,7 @@ def test_zip2dir_round_trip_permissions(tmp_path: Path) -> None:
92169
_write_test_permissions_zip(source_zip)
93170
extract_path = tmp_path / "unzip2"
94171
zip2dir(source_zip, tmp_path / "unzip1")
95-
dir2zip(tmp_path / "unzip1", tmp_path / "tmp.zip")
172+
dir2zip(tmp_path / "unzip1", tmp_path / "tmp.zip", zlib.Z_DEFAULT_COMPRESSION, None)
96173
zip2dir(tmp_path / "tmp.zip", extract_path)
97174
_check_permissions(extract_path)
98175

@@ -104,7 +181,7 @@ def test_dir2zip_deflate(tmp_path: Path) -> None:
104181
input_file = input_dir / "zeros.bin"
105182
input_file.write_bytes(buffer)
106183
output_file = tmp_path / "ouput.zip"
107-
dir2zip(input_dir, output_file)
184+
dir2zip(input_dir, output_file, zlib.Z_DEFAULT_COMPRESSION, None)
108185
assert output_file.stat().st_size < len(buffer) / 4
109186

110187

@@ -117,7 +194,7 @@ def test_dir2zip_folders(tmp_path: Path) -> None:
117194
empty_folder = input_dir / "dummy" / "empty"
118195
empty_folder.mkdir(parents=True)
119196
output_file = tmp_path / "output.zip"
120-
dir2zip(input_dir, output_file)
197+
dir2zip(input_dir, output_file, zlib.Z_DEFAULT_COMPRESSION, None)
121198
expected_dirs = {"dummy/", "dummy/empty/", "dummy-1.0.dist-info/"}
122199
with zipfile.ZipFile(output_file, "r") as z:
123200
assert len(z.filelist) == 4

0 commit comments

Comments
 (0)