forked from STI/Astroport.ONE
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import email.policy
|
|
import itertools
|
|
import os
|
|
from collections.abc import Iterable
|
|
from email.parser import BytesParser
|
|
|
|
from ..wheelfile import WheelFile
|
|
|
|
|
|
def _compute_tags(original_tags: Iterable[str], new_tags: str | None) -> set[str]:
|
|
"""Add or replace tags. Supports dot-separated tags"""
|
|
if new_tags is None:
|
|
return set(original_tags)
|
|
|
|
if new_tags.startswith("+"):
|
|
return {*original_tags, *new_tags[1:].split(".")}
|
|
|
|
if new_tags.startswith("-"):
|
|
return set(original_tags) - set(new_tags[1:].split("."))
|
|
|
|
return set(new_tags.split("."))
|
|
|
|
|
|
def tags(
|
|
wheel: str,
|
|
python_tags: str | None = None,
|
|
abi_tags: str | None = None,
|
|
platform_tags: str | None = None,
|
|
build_tag: str | None = None,
|
|
remove: bool = False,
|
|
) -> str:
|
|
"""Change the tags on a wheel file.
|
|
|
|
The tags are left unchanged if they are not specified. To specify "none",
|
|
use ["none"]. To append to the previous tags, a tag should start with a
|
|
"+". If a tag starts with "-", it will be removed from existing tags.
|
|
Processing is done left to right.
|
|
|
|
:param wheel: The paths to the wheels
|
|
:param python_tags: The Python tags to set
|
|
:param abi_tags: The ABI tags to set
|
|
:param platform_tags: The platform tags to set
|
|
:param build_tag: The build tag to set
|
|
:param remove: Remove the original wheel
|
|
"""
|
|
with WheelFile(wheel, "r") as f:
|
|
assert f.filename, f"{f.filename} must be available"
|
|
|
|
wheel_info = f.read(f.dist_info_path + "/WHEEL")
|
|
info = BytesParser(policy=email.policy.compat32).parsebytes(wheel_info)
|
|
|
|
original_wheel_name = os.path.basename(f.filename)
|
|
namever = f.parsed_filename.group("namever")
|
|
build = f.parsed_filename.group("build")
|
|
original_python_tags = f.parsed_filename.group("pyver").split(".")
|
|
original_abi_tags = f.parsed_filename.group("abi").split(".")
|
|
original_plat_tags = f.parsed_filename.group("plat").split(".")
|
|
|
|
tags: list[str] = info.get_all("Tag", [])
|
|
existing_build_tag = info.get("Build")
|
|
|
|
impls = {tag.split("-")[0] for tag in tags}
|
|
abivers = {tag.split("-")[1] for tag in tags}
|
|
platforms = {tag.split("-")[2] for tag in tags}
|
|
|
|
if impls != set(original_python_tags):
|
|
msg = f"Wheel internal tags {impls!r} != filename tags {original_python_tags!r}"
|
|
raise AssertionError(msg)
|
|
|
|
if abivers != set(original_abi_tags):
|
|
msg = f"Wheel internal tags {abivers!r} != filename tags {original_abi_tags!r}"
|
|
raise AssertionError(msg)
|
|
|
|
if platforms != set(original_plat_tags):
|
|
msg = (
|
|
f"Wheel internal tags {platforms!r} != filename tags {original_plat_tags!r}"
|
|
)
|
|
raise AssertionError(msg)
|
|
|
|
if existing_build_tag != build:
|
|
msg = (
|
|
f"Incorrect filename '{build}' "
|
|
f"& *.dist-info/WHEEL '{existing_build_tag}' build numbers"
|
|
)
|
|
raise AssertionError(msg)
|
|
|
|
# Start changing as needed
|
|
if build_tag is not None:
|
|
build = build_tag
|
|
|
|
final_python_tags = sorted(_compute_tags(original_python_tags, python_tags))
|
|
final_abi_tags = sorted(_compute_tags(original_abi_tags, abi_tags))
|
|
final_plat_tags = sorted(_compute_tags(original_plat_tags, platform_tags))
|
|
|
|
final_tags = [
|
|
namever,
|
|
".".join(final_python_tags),
|
|
".".join(final_abi_tags),
|
|
".".join(final_plat_tags),
|
|
]
|
|
if build:
|
|
final_tags.insert(1, build)
|
|
|
|
final_wheel_name = "-".join(final_tags) + ".whl"
|
|
|
|
if original_wheel_name != final_wheel_name:
|
|
del info["Tag"], info["Build"]
|
|
for a, b, c in itertools.product(
|
|
final_python_tags, final_abi_tags, final_plat_tags
|
|
):
|
|
info["Tag"] = f"{a}-{b}-{c}"
|
|
if build:
|
|
info["Build"] = build
|
|
|
|
original_wheel_path = os.path.join(
|
|
os.path.dirname(f.filename), original_wheel_name
|
|
)
|
|
final_wheel_path = os.path.join(os.path.dirname(f.filename), final_wheel_name)
|
|
|
|
with WheelFile(original_wheel_path, "r") as fin, WheelFile(
|
|
final_wheel_path, "w"
|
|
) as fout:
|
|
fout.comment = fin.comment # preserve the comment
|
|
for item in fin.infolist():
|
|
if item.is_dir():
|
|
continue
|
|
if item.filename == f.dist_info_path + "/RECORD":
|
|
continue
|
|
if item.filename == f.dist_info_path + "/WHEEL":
|
|
fout.writestr(item, info.as_bytes())
|
|
else:
|
|
fout.writestr(item, fin.read(item))
|
|
|
|
if remove:
|
|
os.remove(original_wheel_path)
|
|
|
|
return final_wheel_name
|