feat: add --local-version option to wheel pack

docs/news.rst

@@ -3,6 +3,9 @@ Release Notes
 
 **UNRELEASED**
 
+- Added a ``--local-version`` option to ``wheel pack`` to add, replace, or remove a
+  PEP 440 local version identifier from a wheel
+  (`#570 <https://github.com/pypa/wheel/issues/570>`_)
 - Fixed ``wheel convert`` unnecessarily upgrading compatible core metadata versions
   (`#643 <https://github.com/pypa/wheel/issues/643>`_)
 

pyproject.toml

@@ -67,6 +67,7 @@ exclude = [
     ".github/**",
     ".gitignore",
     ".pre-commit-config.yaml",
+    "**/*.egg-info/**",
     ".readthedocs.yml",
     "**/__pycache__",
 ]

src/wheel/_commands/__init__.py

@@ -21,7 +21,7 @@ def unpack_f(args: argparse.Namespace) -> None:
 def pack_f(args: argparse.Namespace) -> None:
     from .pack import pack
 
-    pack(args.directory, args.dest_dir, args.build_number)
+    pack(args.directory, args.dest_dir, args.build_number, args.local_version)
 
 
 def convert_f(args: argparse.Namespace) -> None:
@@ -104,6 +104,9 @@ def parser() -> argparse.ArgumentParser:
     repack_parser.add_argument(
         "--build-number", help="Build tag to use in the wheel name"
     )
+    repack_parser.add_argument(
+        "--local-version", help="Local version identifier to add or replace"
+    )
     repack_parser.set_defaults(func=pack_f)
 
     convert_parser = s.add_parser("convert", help="Convert egg or wininst to wheel")

src/wheel/_commands/pack.py

@@ -6,19 +6,28 @@
 from email.generator import BytesGenerator
 from email.parser import BytesParser
 
+from packaging.version import InvalidVersion, Version
+
 from ..wheelfile import WheelError, WheelFile
 
 DIST_INFO_RE = re.compile(r"^(?P<namever>(?P<name>.+?)-(?P<ver>\d.*?))\.dist-info$")
 
 
-def pack(directory: str, dest_dir: str, build_number: str | None) -> None:
+def pack(
+    directory: str,
+    dest_dir: str,
+    build_number: str | None,
+    local_version: str | None = None,
+) -> 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)
+    :param build_number: Build tag to use in the wheel name
+    :param local_version: Local version identifier to add or replace
     """
     # Find the .dist-info directory
     dist_info_dirs = [
@@ -33,7 +42,50 @@ def pack(directory: str, dest_dir: str, build_number: str | None) -> None:
 
     # Determine the target wheel filename
     dist_info_dir = dist_info_dirs[0]
-    name_version = DIST_INFO_RE.match(dist_info_dir).group("namever")
+    match = DIST_INFO_RE.match(dist_info_dir)
+    name_version = match.group("namever")
+    name = match.group("name")
+    version = match.group("ver")
+
+    # Update the local version identifier if requested
+    if local_version is not None:
+        base_version = version.split("+", 1)[0]
+        new_version = (
+            base_version if local_version == "" else f"{base_version}+{local_version}"
+        )
+        try:
+            Version(new_version)
+        except InvalidVersion as e:
+            raise WheelError(f"Invalid version {new_version!r}: {e}") from None
+
+        if "-" in new_version:
+            raise WheelError(
+                f"Invalid local version identifier {local_version!r}: "
+                "wheel filenames cannot contain '-' in the version; "
+                "use '_' instead of '-'"
+            )
+
+        new_namever = f"{name}-{new_version}"
+        new_dist_info_dir = f"{new_namever}.dist-info"
+        os.rename(
+            os.path.join(directory, dist_info_dir),
+            os.path.join(directory, new_dist_info_dir),
+        )
+        dist_info_dir = new_dist_info_dir
+        name_version = new_namever
+
+        # Update the Version field in .dist-info/METADATA
+        metadata_path = os.path.join(directory, dist_info_dir, "METADATA")
+        try:
+            with open(metadata_path, "rb") as f:
+                metadata = BytesParser(policy=email.policy.compat32).parse(f)
+        except FileNotFoundError:
+            raise WheelError(f"Missing {dist_info_dir}/METADATA file") from None
+
+        del metadata["Version"]
+        metadata["Version"] = new_version
+        with open(metadata_path, "wb") as f:
+            BytesGenerator(f, maxheaderlen=0).flatten(metadata)
 
     # Read the tags and the existing build number from .dist-info/WHEEL
     wheel_file_path = os.path.join(directory, dist_info_dir, "WHEEL")

tests/commands/test_pack.py

@@ -2,13 +2,18 @@
 
 import email.policy
 import os
+import sys
+from email.generator import BytesGenerator
 from email.message import Message
 from email.parser import BytesParser
+from io import StringIO
 from zipfile import Path, ZipFile
 
 import pytest
 from pytest import TempPathFactory
 
+from wheel._commands import main
+
 from .util import run_command
 
 THISDIR = os.path.dirname(__file__)
@@ -90,3 +95,132 @@ def test_pack(
     assert sorted(new_wheel_file_content.items()) == sorted(
         expected_wheel_content.items()
     )
+
+
+@pytest.mark.parametrize(
+    "local_version_arg, existing_local_version, filename, expected_version",
+    [
+        pytest.param(
+            "l0.cal",
+            None,
+            "test-1.0+l0.cal-py2.py3-none-any.whl",
+            "1.0+l0.cal",
+            id="addlocal",
+        ),
+        pytest.param(
+            "new",
+            "old",
+            "test-1.0+new-py2.py3-none-any.whl",
+            "1.0+new",
+            id="replacelocal",
+        ),
+        pytest.param(
+            "",
+            "old",
+            "test-1.0-py2.py3-none-any.whl",
+            "1.0",
+            id="removelocal",
+        ),
+    ],
+)
+def test_pack_local_version(
+    tmp_path_factory: TempPathFactory,
+    tmp_path: Path,
+    local_version_arg: str,
+    existing_local_version: str | None,
+    filename: str,
+    expected_version: str,
+) -> None:
+    unpack_dir = tmp_path_factory.mktemp("wheeldir")
+    with ZipFile(TESTWHEEL_PATH) as zf:
+        zf.extractall(unpack_dir)
+
+    dist_info_dir = unpack_dir.joinpath("test-1.0.dist-info")
+    if existing_local_version:
+        existing_version = f"1.0+{existing_local_version}"
+        new_dist_info_dir = unpack_dir.joinpath(f"test-{existing_version}.dist-info")
+        dist_info_dir.rename(new_dist_info_dir)
+        dist_info_dir = new_dist_info_dir
+
+        wheel_file_path = dist_info_dir.joinpath("WHEEL")
+        wheel_file_content = wheel_file_path.read_bytes()
+        assert b"Build" not in wheel_file_content
+
+        metadata_path = dist_info_dir.joinpath("METADATA")
+        metadata = BytesParser(policy=email.policy.compat32).parsebytes(
+            metadata_path.read_bytes()
+        )
+        del metadata["Version"]
+        metadata["Version"] = existing_version
+        with open(metadata_path, "wb") as f:
+            BytesGenerator(f, maxheaderlen=0).flatten(metadata)
+
+    args = ["--dest", tmp_path, unpack_dir, "--local-version", local_version_arg]
+    run_command("pack", *args)
+    new_wheel_path = tmp_path.joinpath(filename)
+    assert new_wheel_path.is_file()
+
+    with ZipFile(new_wheel_path) as zf:
+        dist_info_prefix = f"test-{expected_version}.dist-info"
+        parser = BytesParser(policy=email.policy.compat32)
+        metadata = parser.parsebytes(zf.read(f"{dist_info_prefix}/METADATA"))
+        assert metadata["Version"] == expected_version
+
+        wheel_file_content = parser.parsebytes(zf.read(f"{dist_info_prefix}/WHEEL"))
+        assert wheel_file_content["Wheel-Version"] == "1.0"
+
+
+def test_pack_local_version_rejects_hyphen(
+    tmp_path_factory: TempPathFactory, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    unpack_dir = tmp_path_factory.mktemp("wheeldir")
+    with ZipFile(TESTWHEEL_PATH) as zf:
+        zf.extractall(unpack_dir)
+
+    argv = [
+        "wheel",
+        "pack",
+        "--dest",
+        str(tmp_path),
+        str(unpack_dir),
+        "--local-version",
+        "bad-local",
+    ]
+    stdout = StringIO()
+    stderr = StringIO()
+    with monkeypatch.context() as m:
+        m.setattr(sys, "argv", argv)
+        m.setattr(sys, "stdout", stdout)
+        m.setattr(sys, "stderr", stderr)
+        returncode = main()
+
+    assert returncode == 1
+    assert "bad-local" in stderr.getvalue()
+
+
+def test_pack_local_version_rejects_invalid(
+    tmp_path_factory: TempPathFactory, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    unpack_dir = tmp_path_factory.mktemp("wheeldir")
+    with ZipFile(TESTWHEEL_PATH) as zf:
+        zf.extractall(unpack_dir)
+
+    argv = [
+        "wheel",
+        "pack",
+        "--dest",
+        str(tmp_path),
+        str(unpack_dir),
+        "--local-version",
+        "!invalid",
+    ]
+    stdout = StringIO()
+    stderr = StringIO()
+    with monkeypatch.context() as m:
+        m.setattr(sys, "argv", argv)
+        m.setattr(sys, "stdout", stdout)
+        m.setattr(sys, "stderr", stderr)
+        returncode = main()
+
+    assert returncode == 1
+    assert "!invalid" in stderr.getvalue()