refactor: custom impl

Author: henryiii

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: str, build_number: str | None) -> str:
+    """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

@@ -1,4 +1,4 @@
-from __future__ import print_function
+from __future__ import annotations
 
 import itertools
 import os
@@ -8,7 +8,7 @@
 from contextlib import contextmanager
 
 from ..wheelfile import WheelFile
-from .pack import pack
+from .pack import pack, read_tags, set_build_number
 from .unpack import unpack
 
 try:
@@ -84,14 +84,13 @@ def compute_tags(original_tags, new_tags):
 
 
 def tags(
-    wheels,
-    python_tags=None,
-    abi_tags=None,
-    platform_tags=None,
-    build_number=None,
-    remove=False,
-):
-    # type: (List[str], Optional[List[str]], Optional[List[str]], Optional[List[str]], Optional[int], bool) -> Iterator[str]  # noqa: E501
+    wheels: list[str],
+    python_tags: list[str] | None = None,
+    abi_tags: list[str] | None = None,
+    platform_tags: list[str] = None,
+    build_number: int | None = None,
+    remove: bool = False,
+) -> Iterator[str]:
     """Change the tags on a wheel file.
 
     The tags are left unchanged if they are not specified. To specify "none",
@@ -106,54 +105,103 @@ def tags(
     """
 
     for wheel in wheels:
-        with temporary_directory() as tmpdir, InWheelCtx(wheel, tmpdir) as wfctx:
-            namever = wfctx.parsed_filename.group("namever")
-            build = wfctx.parsed_filename.group("build")
-            original_python_tags = wfctx.parsed_filename.group("pyver").split(".")
-            original_abi_tags = wfctx.parsed_filename.group("abi").split(".")
-            orignial_plat_tags = wfctx.parsed_filename.group("plat").split(".")
-
-            if build_number is not None:
-                build = str(build_number)
-
-            final_python_tags = compute_tags(original_python_tags, python_tags)
-            final_abi_tags = compute_tags(original_abi_tags, abi_tags)
-            final_plat_tags = compute_tags(orignial_plat_tags, platform_tags)
-
-            final_tags = [
-                ".".join(sorted(final_python_tags)),
-                ".".join(sorted(final_abi_tags)),
-                ".".join(sorted(final_plat_tags)),
+        with WheelFile(wheel, "r") as f:
+            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(".")
+            orignial_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):
+            raise AssertionError(f"{impls} != {original_python_tags}")
+
+        if abivers != set(original_abi_tags):
+            raise AssertionError(f"{abivers} != {original_abi_tags}")
+
+        if platforms != set(orignial_plat_tags):
+            raise AssertionError(f"{platforms} != {orignial_plat_tags}")
+
+        if existing_build_number != build:
+            raise AssertionError(
+                f"Incorrect filename '{build}' & "
+                f"*.dist-info/WHEEL '{existing_build_number}' build numbers"
+            )
+
+        # Start changing as needed
+        if build_number is not None:
+            build = str(build_number)
+
+        final_python_tags = compute_tags(original_python_tags, python_tags)
+        final_abi_tags = compute_tags(original_abi_tags, abi_tags)
+        final_plat_tags = compute_tags(orignial_plat_tags, platform_tags)
+
+        final_tags = [
+            ".".join(sorted(final_python_tags)),
+            ".".join(sorted(final_abi_tags)),
+            ".".join(sorted(final_plat_tags)),
+        ]
+
+        if build:
+            final_tags.insert(0, build)
+        final_tags.insert(0, namever)
+
+        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
+                )
             ]
 
-            if build:
-                final_tags.insert(0, build)
-            final_tags.insert(0, namever)
+            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)
 
-            original_wheel_name = os.path.basename(wfctx.filename)
-            final_wheel_name = "-".join(final_tags) + ".whl"
+        yield final_wheel_name
 
-            if original_wheel_name != final_wheel_name:
 
-                wheelinfo = os.path.join(
-                    tmpdir, namever, wfctx.wheel.dist_info_path, "WHEEL"
-                )
-                with open(wheelinfo, "rb+") as f:
-                    lines = [line for line in f if not line.startswith(b"Tag:")]
-                    for a, b, c in itertools.product(
-                        final_python_tags, final_abi_tags, final_plat_tags
-                    ):
-                        lines.append(
-                            "Tag: {}-{}-{}\r\n".format(a, b, c).encode("ascii")
-                        )
-                    f.seek(0)
-                    f.truncate()
-                    f.write(b"".join(lines))
-
-                wfctx.build_number = build
-                wfctx.dirname = os.path.dirname(wheel)
-
-        if remove:
-            os.remove(wheel)
+def set_tags(in_string: bytes, tags: list[str]) -> bytes:
+    """Set the tags in the .dist-info/WHEEL file contents.
 
-        yield final_wheel_name
+    :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

tests/cli/test_tags.py

@@ -1,4 +1,5 @@
 import os
+from zipfile import ZipFile
 
 import py
 import pytest
@@ -166,3 +167,31 @@ def test_tags_command_del(capsys, wheelpath):
     assert "test-1.0-py2.py3.py4-cp33m-linux_x86_64.whl" == newname
     output_file = wheelpath.dirpath(newname)
     output_file.remove()
+
+
+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.dirpath(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.remove()