forked from bazel-contrib/rules_python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrepack_whl.py
More file actions
182 lines (145 loc) · 5.66 KB
/
repack_whl.py
File metadata and controls
182 lines (145 loc) · 5.66 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
# Copyright 2023 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Regenerate a whl file after patching and cleanup the patched contents.
This script will take contents of the current directory and create a new wheel
out of it and will remove all files that were written to the wheel.
"""
from __future__ import annotations
import argparse
import difflib
import logging
import pathlib
import sys
import tempfile
from tools.wheelmaker import _WhlFile
# NOTE: Implement the following matching of what goes into the RECORD
# https://peps.python.org/pep-0491/#the-dist-info-directory
_EXCLUDES = [
"RECORD",
"INSTALLER",
"RECORD.jws",
"RECORD.p7s",
"REQUESTED",
]
_DISTINFO = "dist-info"
def _unidiff_output(expected, actual, record):
"""
Helper function. Returns a string containing the unified diff of two
multiline strings.
"""
expected = expected.splitlines(1)
actual = actual.splitlines(1)
diff = difflib.unified_diff(
expected, actual, fromfile=f"a/{record}", tofile=f"b/{record}"
)
return "".join(diff)
def _files_to_pack(dir: pathlib.Path, want_record: str) -> list[pathlib.Path]:
"""Check that the RECORD file entries are correct and print a unified diff on failure."""
# First get existing files by using the RECORD file
got_files = []
got_distinfos = []
for line in want_record.splitlines():
rec, _, _ = line.partition(",")
path = dir / rec
if not path.exists():
# skip files that do not exist as they won't be present in the final
# RECORD file.
continue
if not path.parent.name.endswith(_DISTINFO):
got_files.append(path)
elif path.name not in _EXCLUDES:
got_distinfos.append(path)
# Then get extra files present in the directory but not in the RECORD file
extra_files = []
extra_distinfos = []
for path in dir.rglob("*"):
if path.is_dir():
continue
elif path.parent.name.endswith(_DISTINFO):
if path.name in _EXCLUDES:
# NOTE: we implement the following matching of what goes into the RECORD
# https://peps.python.org/pep-0491/#the-dist-info-directory
continue
elif path not in got_distinfos:
extra_distinfos.append(path)
elif path not in got_files:
extra_files.append(path)
# sort the extra files for reproducibility
extra_files.sort()
extra_distinfos.sort()
# This order ensures that the structure of the RECORD file is always the
# same and ensures smaller patchsets to the RECORD file in general
return got_files + extra_files + got_distinfos + extra_distinfos
def main(sys_argv):
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"whl_path",
type=pathlib.Path,
help="The original wheel file that we have patched.",
)
parser.add_argument(
"--record-patch",
type=pathlib.Path,
help="The output path that we are going to write the RECORD file patch to.",
)
parser.add_argument(
"output",
type=pathlib.Path,
help="The output path that we are going to write a new file to.",
)
args = parser.parse_args(sys_argv)
cwd = pathlib.Path.cwd()
logging.debug("=" * 80)
logging.debug("Repackaging the wheel")
logging.debug("=" * 80)
with tempfile.TemporaryDirectory(dir=cwd) as tmpdir:
patched_wheel_dir = cwd / tmpdir
logging.debug(f"Created a tmpdir: {patched_wheel_dir}")
excludes = [args.whl_path, patched_wheel_dir]
logging.debug("Moving whl contents to the newly created tmpdir")
for p in cwd.glob("*"):
if p in excludes:
logging.debug(f"Ignoring: {p}")
continue
rel_path = p.relative_to(cwd)
dst = p.rename(patched_wheel_dir / rel_path)
logging.debug(f"mv {p} -> {dst}")
distinfo_dir = next(iter(patched_wheel_dir.glob("*dist-info")))
logging.debug(f"Found dist-info dir: {distinfo_dir}")
record_path = distinfo_dir / "RECORD"
record_contents = record_path.read_text() if record_path.exists() else ""
with _WhlFile(args.output, mode="w", distinfo_dir=distinfo_dir) as out:
for p in _files_to_pack(patched_wheel_dir, record_contents):
rel_path = p.relative_to(patched_wheel_dir)
out.add_file(str(rel_path), p)
logging.debug(f"Writing RECORD file")
got_record = out.add_recordfile().decode("utf-8", "surrogateescape")
if got_record == record_contents:
logging.info(f"Created a whl file: {args.output}")
return
record_diff = _unidiff_output(
record_contents,
got_record,
out.distinfo_path("RECORD"),
)
args.record_patch.write_text(record_diff)
logging.warning(
f"Please apply patch to the RECORD file ({args.record_patch}):\n{record_diff}"
)
if __name__ == "__main__":
logging.basicConfig(
format="%(module)s: %(levelname)s: %(message)s", level=logging.DEBUG
)
sys.exit(main(sys.argv[1:]))