Added tags CLI interface (#422)

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

6 files changed, 532 insertions, 21 deletions

diff --git a/docs/reference/index.rst b/docs/reference/index.rst
index 1921323..f332026 100644
--- a/docs/reference/index.rst
+++ b/docs/reference/index.rst
@@ -7,3 +7,4 @@ Reference Guide
    wheel_convert
    wheel_unpack
    wheel_pack
+   wheel_tags
diff --git a/docs/reference/wheel_tags.rst b/docs/reference/wheel_tags.rst
new file mode 100644
index 0000000..f437071
--- /dev/null
+++ b/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
diff --git a/src/wheel/cli/__init__.py b/src/wheel/cli/__init__.py
index c0fb8c4..fa1f10b 100644
--- a/src/wheel/cli/__init__.py
+++ b/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)
 
diff --git a/src/wheel/cli/pack.py b/src/wheel/cli/pack.py
index 1949d4c..7c75c63 100644
--- a/src/wheel/cli/pack.py
+++ b/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)])
diff --git a/src/wheel/cli/tags.py b/src/wheel/cli/tags.py
new file mode 100644
index 0000000..0ea0f44
--- /dev/null
+++ b/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
diff --git a/tests/cli/test_tags.py b/tests/cli/test_tags.py
new file mode 100644
index 0000000..630c903
--- /dev/null
+++ b/tests/cli/test_tags.py
@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from zipfile import ZipFile
+
+import pytest
+
+from wheel.cli import parser
+from wheel.cli.tags import tags
+from wheel.wheelfile import WheelFile
+
+TESTDIR = Path(__file__).parent.parent
+TESTWHEEL_NAME = "test-1.0-py2.py3-none-any.whl"
+TESTWHEEL_PATH = TESTDIR / "testdata" / TESTWHEEL_NAME
+
+
+@pytest.fixture
+def wheelpath(tmp_path):
+    wheels_dir = tmp_path / "wheels"
+    wheels_dir.mkdir()
+    fn = wheels_dir / TESTWHEEL_NAME
+    # The str calls can be removed for Python 3.8+
+    shutil.copy(str(TESTWHEEL_PATH), str(fn))
+    return fn
+
+
+def test_tags_no_args(wheelpath):
+    newname = tags(str(wheelpath))
+    assert TESTWHEEL_NAME == newname
+    assert wheelpath.exists()
+
+
+def test_python_tags(wheelpath):
+    newname = tags(str(wheelpath), python_tags="py3")
+    assert TESTWHEEL_NAME.replace("py2.py3", "py3") == newname
+    output_file = wheelpath.parent / newname
+    with WheelFile(str(output_file)) as f:
+        output = f.read(f.dist_info_path + "/WHEEL")
+    assert (
+        output == b"Wheel-Version: 1.0\r\nGenerator: bdist_wheel (0.30.0)"
+        b"\r\nRoot-Is-Purelib: false\r\nTag: py3-none-any\r\n"
+    )
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), python_tags="py2.py3")
+    assert TESTWHEEL_NAME == newname
+
+    newname = tags(str(wheelpath), python_tags="+py4", remove=True)
+    assert not wheelpath.exists()
+    assert TESTWHEEL_NAME.replace("py2.py3", "py2.py3.py4") == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+
+def test_abi_tags(wheelpath):
+    newname = tags(str(wheelpath), abi_tags="cp33m")
+    assert TESTWHEEL_NAME.replace("none", "cp33m") == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), abi_tags="cp33m.abi3")
+    assert TESTWHEEL_NAME.replace("none", "abi3.cp33m") == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), abi_tags="none")
+    assert TESTWHEEL_NAME == newname
+
+    newname = tags(str(wheelpath), abi_tags="+abi3.cp33m", remove=True)
+    assert not wheelpath.exists()
+    assert TESTWHEEL_NAME.replace("none", "abi3.cp33m.none") == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+
+def test_plat_tags(wheelpath):
+    newname = tags(str(wheelpath), platform_tags="linux_x86_64")
+    assert TESTWHEEL_NAME.replace("any", "linux_x86_64") == newname
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), platform_tags="linux_x86_64.win32")
+    assert TESTWHEEL_NAME.replace("any", "linux_x86_64.win32") == newname
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), platform_tags="+linux_x86_64.win32")
+    assert TESTWHEEL_NAME.replace("any", "any.linux_x86_64.win32") == newname
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+    output_file.unlink()
+
+    newname = tags(str(wheelpath), platform_tags="+linux_x86_64.win32")
+    assert TESTWHEEL_NAME.replace("any", "any.linux_x86_64.win32") == newname
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+
+    newname2 = tags(str(output_file), platform_tags="-any")
+    output_file.unlink()
+
+    assert TESTWHEEL_NAME.replace("any", "linux_x86_64.win32") == newname2
+    output_file2 = wheelpath.parent / newname2
+    assert output_file2.exists()
+    output_file2.unlink()
+
+    newname = tags(str(wheelpath), platform_tags="any")
+    assert TESTWHEEL_NAME == newname
+
+
+def test_build_number(wheelpath):
+    newname = tags(str(wheelpath), build_number=1)
+    assert TESTWHEEL_NAME.replace("-py2", "-1-py2") == newname
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+    output_file.unlink()
+
+
+def test_multi_tags(wheelpath):
+    newname = tags(
+        str(wheelpath),
+        platform_tags="linux_x86_64",
+        python_tags="+py4",
+        build_number=1,
+    )
+    assert "test-1.0-1-py2.py3.py4-none-linux_x86_64.whl" == newname
+
+    output_file = wheelpath.parent / newname
+    assert output_file.exists()
+    with WheelFile(str(output_file)) as f:
+        output = f.read(f.dist_info_path + "/WHEEL")
+    assert (
+        output
+        == b"Wheel-Version: 1.0\r\nGenerator: bdist_wheel (0.30.0)\r\nRoot-Is-Purelib:"
+        b" false\r\nTag: py2-none-linux_x86_64\r\nTag: py3-none-linux_x86_64\r\nTag:"
+        b" py4-none-linux_x86_64\r\nBuild: 1\r\n"
+    )
+    output_file.unlink()
+
+
+def test_tags_command(capsys, wheelpath):
+    args = [
+        "tags",
+        "--python-tag",
+        "py3",
+        "--abi-tag",
+        "cp33m",
+        "--platform-tag",
+        "linux_x86_64",
+        "--build",
+        "7",
+        str(wheelpath),
+    ]
+    p = parser()
+    args = p.parse_args(args)
+    args.func(args)
+    assert wheelpath.exists()
+
+    newname = capsys.readouterr().out.strip()
+    assert "test-1.0-7-py3-cp33m-linux_x86_64.whl" == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+
+def test_tags_command_del(capsys, wheelpath):
+    args = [
+        "tags",
+        "--python-tag",
+        "+py4",
+        "--abi-tag",
+        "cp33m",
+        "--platform-tag",
+        "linux_x86_64",
+        "--remove",
+        str(wheelpath),
+    ]
+    p = parser()
+    args = p.parse_args(args)
+    args.func(args)
+    assert not wheelpath.exists()
+
+    newname = capsys.readouterr().out.strip()
+    assert "test-1.0-py2.py3.py4-cp33m-linux_x86_64.whl" == newname
+    output_file = wheelpath.parent / newname
+    output_file.unlink()
+
+
+def test_permission_bits(capsys, wheelpath):
+    args = [
+        "tags",
+        "--python-tag=+py4",
+        str(wheelpath),
+    ]
+    p = parser()
+    args = p.parse_args(args)
+    args.func(args)
+
+    newname = capsys.readouterr().out.strip()
+    assert "test-1.0-py2.py3.py4-none-any.whl" == newname
+    output_file = wheelpath.parent / newname
+
+    with ZipFile(str(output_file), "r") as outf:
+        with ZipFile(str(wheelpath), "r") as inf:
+            for member in inf.namelist():
+                if not member.endswith("/RECORD"):
+                    out_attr = outf.getinfo(member).external_attr
+                    inf_attr = inf.getinfo(member).external_attr
+                    assert (
+                        out_attr == inf_attr
+                    ), f"{member} 0x{out_attr:012o} != 0x{inf_attr:012o}"
+
+    output_file.unlink()