Added tags CLI interface (#422)

Co-authored-by: Joe Rickerby <joerick@mac.com>

docs/reference/index.rst

@@ -7,3 +7,4 @@ Reference Guide
    wheel_convert
    wheel_unpack
    wheel_pack
+   wheel_tags

docs/reference/wheel_tags.rst

@@ -0,0 +1,62 @@
+wheel tags
+==========
+
+Usage
+-----
+
+::
+
+    wheel tags [-h] [--remove] [--python-tag TAG] [--abi-tag TAG] [--platform-tag TAG] [--build NUMBER] WHEEL [...]
+
+Description
+-----------
+
+Make a new wheel with given tags from and existing wheel. Any tags left
+unspecified will remain the same. Multiple tags are separated by a "." Starting
+with a "+" will append to the existing tags.  Starting with a "-" will remove a
+tag. Be sure to use the equals syntax on the shell so that it does not get
+parsed as an extra option, such as ``--python-tag=-py2``. The original file
+will remain unless ``--remove`` is given. The output filename(s) will be
+displayed on stdout for further processing.
+
+
+Options
+-------
+
+.. option:: --remove
+
+    Remove the original wheel, keeping only the retagged wheel.
+
+.. option:: --python-tag=TAG
+
+    Override the python tag (prepend with "+" to append, "-" to remove).
+    Multiple tags can be separated with a dot.
+
+.. option:: --abi-tag=TAG
+
+    Override the abi tag (prepend with "+" to append, "-" to remove).
+    Multiple tags can be separated with a dot.
+
+.. option:: --platform-tag=TAG
+
+    Override the platform tag (prepend with "+" to append, "-" to remove).
+    Multiple tags can be separated with a dot.
+
+.. option:: --build=NUMBER
+
+    Specify a build number.
+
+Examples
+--------
+
+* Replace a wheel's Python specific tags with generic tags (if no Python extensions are present, for example)::
+
+    $ wheel tags --python-tag=py2.py3 --abi-tag=none cmake-3.20.2-cp39-cp39-win_amd64.whl
+    cmake-3.20.2-py2.py3-none-win_amd64.whl
+
+* Add compatibility tags for macOS universal wheels and older pips::
+
+    $ wheel tags \
+        --platform-tag=+macosx_10_9_x86_64.macosx_11_0_arm64 \
+        ninja-1.11.1-py2.py3-none-macosx_10_9_universal2.whl
+    ninja-1.11.1-py2.py3-none-macosx_10_9_universal2.macosx_10_9_x86_64.macosx_11_0_arm64.whl

src/wheel/cli/__init__.py

@@ -31,12 +31,40 @@ def convert_f(args):
     convert(args.files, args.dest_dir, args.verbose)
 
 
+def tags_f(args):
+    from .tags import tags
+
+    names = (
+        tags(
+            wheel,
+            args.python_tag,
+            args.abi_tag,
+            args.platform_tag,
+            args.build,
+            args.remove,
+        )
+        for wheel in args.wheel
+    )
+
+    for name in names:
+        print(name)
+
+
 def version_f(args):
     from .. import __version__
 
     print("wheel %s" % __version__)
 
 
+TAGS_HELP = """\
+Make a new wheel with given tags. Any tags unspecified will remain the same.
+Starting the tags with a "+" will append to the existing tags. Starting with a
+"-" will remove a tag (use --option=-TAG syntax). Multiple tags can be
+separated by ".". The original file will remain unless --remove is given.  The
+output filename(s) will be displayed on stdout for further processing.
+"""
+
+
 def parser():
     p = argparse.ArgumentParser()
     s = p.add_subparsers(help="commands")
@@ -72,6 +100,27 @@ def parser():
     convert_parser.add_argument("--verbose", "-v", action="store_true")
     convert_parser.set_defaults(func=convert_f)
 
+    tags_parser = s.add_parser(
+        "tags", help="Add or replace the tags on a wheel", description=TAGS_HELP
+    )
+    tags_parser.add_argument("wheel", nargs="*", help="Existing wheel(s) to retag")
+    tags_parser.add_argument(
+        "--remove",
+        action="store_true",
+        help="Remove the original files, keeping only the renamed ones",
+    )
+    tags_parser.add_argument(
+        "--python-tag", metavar="TAG", help="Specify an interpreter tag(s)"
+    )
+    tags_parser.add_argument("--abi-tag", metavar="TAG", help="Specify an ABI tag(s)")
+    tags_parser.add_argument(
+        "--platform-tag", metavar="TAG", help="Specify a platform tag(s)"
+    )
+    tags_parser.add_argument(
+        "--build", type=int, metavar="NUMBER", help="Specify a build number"
+    )
+    tags_parser.set_defaults(func=tags_f)
+
     version_parser = s.add_parser("version", help="Print version and exit")
     version_parser.set_defaults(func=version_f)
 

src/wheel/cli/pack.py

@@ -37,13 +37,8 @@ def pack(directory: str, dest_dir: str, build_number: str | None):
     # 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) as f:
-        tags = []
-        for line in f:
-            if line.startswith("Tag: "):
-                tags.append(line.split(" ")[1].rstrip())
-            elif line.startswith("Build: "):
-                existing_build_number = line.split(" ")[1].rstrip()
+    with open(wheel_file_path, "rb") as f:
+        tags, existing_build_number = read_tags(f.read())
 
         if not tags:
             raise WheelError(
@@ -58,28 +53,16 @@ def pack(directory: str, dest_dir: str, build_number: str | None):
             name_version += "-" + build_number
 
         if build_number != existing_build_number:
-            replacement = (
-                ("Build: %s\r\n" % build_number).encode("ascii")
-                if build_number
-                else b""
-            )
             with open(wheel_file_path, "rb+") as f:
                 wheel_file_content = f.read()
-                wheel_file_content, num_replaced = BUILD_NUM_RE.subn(
-                    replacement, wheel_file_content
-                )
-                if not num_replaced:
-                    wheel_file_content += replacement
+                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
-    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})
-    tagline = "-".join([".".join(impls), ".".join(abivers), ".".join(platforms)])
+    tagline = compute_tagline(tags)
 
     # Repack the wheel
     wheel_path = os.path.join(dest_dir, f"{name_version}-{tagline}.whl")
@@ -88,3 +71,54 @@ def pack(directory: str, dest_dir: str, build_number: str | None):
         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)])

src/wheel/cli/tags.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import itertools
+import os
+from collections.abc import Iterable
+
+from ..wheelfile import WheelFile
+from .pack import read_tags, set_build_number
+
+
+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_number: int | 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_number: The build number 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")
+
+        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, existing_build_number = read_tags(wheel_info)
+
+    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_number != build:
+        msg = (
+            f"Incorrect filename '{build}' "
+            "& *.dist-info/WHEEL '{existing_build_number}' build numbers"
+        )
+        raise AssertionError(msg)
+
+    # Start changing as needed
+    if build_number is not None:
+        build = str(build_number)
+
+    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:
+        tags = [
+            f"{a}-{b}-{c}"
+            for a, b, c in itertools.product(
+                final_python_tags, final_abi_tags, final_plat_tags
+            )
+        ]
+
+        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.filename == f.dist_info_path + "/RECORD":
+                    continue
+                if item.filename == f.dist_info_path + "/WHEEL":
+                    content = fin.read(item)
+                    content = set_tags(content, tags)
+                    content = set_build_number(content, build)
+                    fout.writestr(item, content)
+                else:
+                    fout.writestr(item, fin.read(item))
+
+        if remove:
+            os.remove(original_wheel_path)
+
+    return final_wheel_name
+
+
+def set_tags(in_string: bytes, tags: Iterable[str]) -> bytes:
+    """Set the tags in the .dist-info/WHEEL file contents.
+
+    :param in_string: The string to modify.
+    :param tags: The tags to set.
+    """
+
+    lines = [line for line in in_string.splitlines() if not line.startswith(b"Tag:")]
+    for tag in tags:
+        lines.append(b"Tag: " + tag.encode("ascii"))
+    in_string = b"\r\n".join(lines) + b"\r\n"
+
+    return in_string