Spaces:
Running
Running
| from __future__ import annotations | |
| import os.path | |
| import re | |
| from wheel.cli import WheelError | |
| from wheel.wheelfile import WheelFile | |
| DIST_INFO_RE = re.compile(r"^(?P<namever>(?P<name>.+?)-(?P<ver>\d.*?))\.dist-info$") | |
| BUILD_NUM_RE = re.compile(rb"Build: (\d\w*)$") | |
| def pack(directory: str, dest_dir: str, build_number: str | None): | |
| """Repack a previously unpacked wheel directory into a new wheel file. | |
| The .dist-info/WHEEL file must contain one or more tags so that the target | |
| wheel file name can be determined. | |
| :param directory: The unpacked wheel directory | |
| :param dest_dir: Destination directory (defaults to the current directory) | |
| """ | |
| # Find the .dist-info directory | |
| dist_info_dirs = [ | |
| fn | |
| for fn in os.listdir(directory) | |
| if os.path.isdir(os.path.join(directory, fn)) and DIST_INFO_RE.match(fn) | |
| ] | |
| if len(dist_info_dirs) > 1: | |
| raise WheelError(f"Multiple .dist-info directories found in {directory}") | |
| elif not dist_info_dirs: | |
| raise WheelError(f"No .dist-info directories found in {directory}") | |
| # Determine the target wheel filename | |
| dist_info_dir = dist_info_dirs[0] | |
| name_version = DIST_INFO_RE.match(dist_info_dir).group("namever") | |
| # Read the tags and the existing build number from .dist-info/WHEEL | |
| existing_build_number = None | |
| wheel_file_path = os.path.join(directory, dist_info_dir, "WHEEL") | |
| with open(wheel_file_path, "rb") as f: | |
| tags, existing_build_number = read_tags(f.read()) | |
| if not tags: | |
| raise WheelError( | |
| "No tags present in {}/WHEEL; cannot determine target wheel " | |
| "filename".format(dist_info_dir) | |
| ) | |
| # Set the wheel file name and add/replace/remove the Build tag in .dist-info/WHEEL | |
| build_number = build_number if build_number is not None else existing_build_number | |
| if build_number is not None: | |
| if build_number: | |
| name_version += "-" + build_number | |
| if build_number != existing_build_number: | |
| with open(wheel_file_path, "rb+") as f: | |
| wheel_file_content = f.read() | |
| wheel_file_content = set_build_number(wheel_file_content, build_number) | |
| f.seek(0) | |
| f.truncate() | |
| f.write(wheel_file_content) | |
| # Reassemble the tags for the wheel file | |
| tagline = compute_tagline(tags) | |
| # Repack the wheel | |
| wheel_path = os.path.join(dest_dir, f"{name_version}-{tagline}.whl") | |
| with WheelFile(wheel_path, "w") as wf: | |
| print(f"Repacking wheel as {wheel_path}...", end="", flush=True) | |
| wf.write_files(directory) | |
| print("OK") | |
| def read_tags(input_str: bytes) -> tuple[list[str], str | None]: | |
| """Read tags from a string. | |
| :param input_str: A string containing one or more tags, separated by spaces | |
| :return: A list of tags and a list of build tags | |
| """ | |
| tags = [] | |
| existing_build_number = None | |
| for line in input_str.splitlines(): | |
| if line.startswith(b"Tag: "): | |
| tags.append(line.split(b" ")[1].rstrip().decode("ascii")) | |
| elif line.startswith(b"Build: "): | |
| existing_build_number = line.split(b" ")[1].rstrip().decode("ascii") | |
| return tags, existing_build_number | |
| def set_build_number(wheel_file_content: bytes, build_number: str | None) -> bytes: | |
| """Compute a build tag and add/replace/remove as necessary. | |
| :param wheel_file_content: The contents of .dist-info/WHEEL | |
| :param build_number: The build tags present in .dist-info/WHEEL | |
| :return: The (modified) contents of .dist-info/WHEEL | |
| """ | |
| replacement = ( | |
| ("Build: %s\r\n" % build_number).encode("ascii") if build_number else b"" | |
| ) | |
| wheel_file_content, num_replaced = BUILD_NUM_RE.subn( | |
| replacement, wheel_file_content | |
| ) | |
| if not num_replaced: | |
| wheel_file_content += replacement | |
| return wheel_file_content | |
| def compute_tagline(tags: list[str]) -> str: | |
| """Compute a tagline from a list of tags. | |
| :param tags: A list of tags | |
| :return: A tagline | |
| """ | |
| impls = sorted({tag.split("-")[0] for tag in tags}) | |
| abivers = sorted({tag.split("-")[1] for tag in tags}) | |
| platforms = sorted({tag.split("-")[2] for tag in tags}) | |
| return "-".join([".".join(impls), ".".join(abivers), ".".join(platforms)]) | |