feat: adding tags CLI interface

henryiii · henryiii · commit 5fdc93710a49 · 2021-10-07T23:07:55.000-04:00
usage: wheel tags [-h] [--remove] [--python-tag TAG] [--abi-tag TAG] [--platform-tag TAG] [--build NUMBER] [wheel ...]

Make a new wheel with given tags. Any tags unspecified will remain the
same. Separate multiple tags with a dot. Starting with a dot will append
to the existing tags. The original file will remain unless --remove is
given. The output file(s) will be displayed on stdout.

positional arguments:
  wheel               Existing wheel(s) to retag

options:
  -h, --help          Show this help message and exit
  --remove            Remove the original files, keeping only the renamed ones
  --python-tag TAG    Specify an interpreter tag(s)
  --abi-tag TAG       Specify an ABI tag(s)
  --platform-tag TAG  Specify a platform tag(s)
  --build NUMBER      Specify a build number
diff --git a/src/wheel/cli/__init__.py b/src/wheel/cli/__init__.py
@@ -35,11 +35,33 @@ def convert_f(args):
     convert(args.files, args.dest_dir, args.verbose)
 
 
+def tags_f(args):
+    from .tags import tags
+
+    for name in tags(
+        args.wheel,
+        args.python_tag and args.python_tag.split('.'),
+        args.abi_tag and args.abi_tag.split('.'),
+        args.platform_tag and args.platform_tag.split('.'),
+        args.build,
+        args.remove
+    ):
+        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.
+Separate multiple tags with a dot. Starting with a dot will append to the
+existing tags. The original file will remain unless --remove is given. The
+output file(s) will be displayed on stdout.
+'''
+
+
 def parser():
     p = argparse.ArgumentParser()
     s = p.add_subparsers(help="commands")
@@ -64,6 +86,17 @@ 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/tags.py b/src/wheel/cli/tags.py
@@ -0,0 +1,157 @@
+from __future__ import print_function
+
+import os
+import sys
+import tempfile
+import itertools
+from contextlib import contextmanager
+import shutil
+
+from ..wheelfile import WheelFile
+from .pack import pack
+from .unpack import unpack
+
+try:
+    from typing import List, Optional, Iterator
+except ImportError:
+    pass
+
+
+@contextmanager
+def redirect_stdout(new_target):
+    old_target, sys.stdout = sys.stdout, new_target
+    try:
+        yield new_target
+    finally:
+        sys.stdout = old_target
+
+
+@contextmanager
+def temporary_directory():
+    try:
+        dirname = tempfile.mkdtemp()
+        yield dirname
+    finally:
+        shutil.rmtree(dirname)
+
+
+class InWheelCtx(object):
+    @property
+    def parsed_filename(self):
+        return self.wheel.parsed_filename
+
+    @property
+    def filename(self):
+        return self.wheel.filename
+
+    def __init__(self, wheel, tmpdir):
+        self.wheel = WheelFile(wheel)
+        self.tmpdir = tmpdir
+        self.build_number = None
+        # If dirname is unset, don't pack a new wheel
+        self.dirname = None
+
+    def __enter__(self):
+        with redirect_stdout(sys.stderr):
+            unpack(self.wheel.filename, self.tmpdir)
+        return self
+
+    def __exit__(self, *args):
+        if self.dirname:
+            with redirect_stdout(sys.stderr):
+                pack(
+                    os.path.join(
+                        self.tmpdir, self.wheel.parsed_filename.group("namever")
+                    ),
+                    self.dirname,
+                    self.build_number,
+                )
+
+
+def compute_tags(original_tags, new_tags):
+    # type: (List[str], Optional[List[str]]) -> List[str]
+    """Add or replace tags."""
+
+    if not new_tags:
+        return original_tags
+
+    if new_tags[0] == '':
+        return original_tags + new_tags[1:]
+    else:
+        return 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
+    """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, use ["", ...].
+
+    :param wheels: 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.
+    """
+
+    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)),
+            ]
+
+            if build:
+                final_tags.insert(0, build)
+            final_tags.insert(0, namever)
+
+            original_wheel_name = os.path.basename(wfctx.filename)
+            final_wheel_name = '-'.join(final_tags) + '.whl'
+
+            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)
+
+        yield final_wheel_name
diff --git a/tests/cli/test_tags.py b/tests/cli/test_tags.py
@@ -0,0 +1,161 @@
+from wheel.cli.tags import tags
+import os
+import py
+import pytest
+
+from wheel.wheelfile import WheelFile
+from wheel.cli import parser
+
+THISDIR = os.path.dirname(__file__)
+TESTWHEEL_NAME = 'test-1.0-py2.py3-none-any.whl'
+TESTWHEEL_PATH = os.path.join(THISDIR, '..', 'testdata', TESTWHEEL_NAME)
+
+
+@pytest.fixture
+def wheelpath(tmpdir):
+    fn = tmpdir.mkdir("wheels").join(TESTWHEEL_NAME)
+    py.path.local(TESTWHEEL_PATH).copy(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.dirpath(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.remove()
+
+    (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.dirpath(newname)
+    output_file.remove()
+
+
+def test_abi_tags(wheelpath):
+    (newname,) = tags([str(wheelpath)], abi_tags=['cp33m'])
+    assert TESTWHEEL_NAME.replace('none', 'cp33m') == newname
+    output_file = wheelpath.dirpath(newname)
+    output_file.remove()
+
+    (newname,) = tags([str(wheelpath)], abi_tags=['abi3', 'cp33m'])
+    assert TESTWHEEL_NAME.replace('none', 'abi3.cp33m') == newname
+    output_file = wheelpath.dirpath(newname)
+    output_file.remove()
+
+    (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.dirpath(newname)
+    output_file.remove()
+
+
+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.dirpath(newname)
+    assert output_file.exists()
+    output_file.remove()
+
+    (newname,) = tags([str(wheelpath)], platform_tags=['linux_x86_64', 'win32'])
+    assert TESTWHEEL_NAME.replace('any', 'linux_x86_64.win32') == newname
+    output_file = wheelpath.dirpath(newname)
+    assert output_file.exists()
+    output_file.remove()
+
+    (newname,) = tags([str(wheelpath)], platform_tags=['', 'linux_x86_64', 'win32'])
+    assert TESTWHEEL_NAME.replace('any', 'any.linux_x86_64.win32') == newname
+    output_file = wheelpath.dirpath(newname)
+    assert output_file.exists()
+    output_file.remove()
+
+    (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.dirpath(newname)
+    assert output_file.exists()
+    output_file.remove()
+
+
+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.dirpath(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.remove()
+
+
+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.dirpath(newname)
+    output_file.remove()
+
+
+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.dirpath(newname)
+    output_file.remove()
