Add compatible tags selector utility

Author: sbidoul

src/packaging/tags.py

@@ -13,16 +13,23 @@
 import sysconfig
 from importlib.machinery import EXTENSION_SUFFIXES
 from typing import (
+    TYPE_CHECKING,
     Any,
     Iterable,
     Iterator,
     Sequence,
     Tuple,
+    TypeVar,
     cast,
 )
 
 from . import _manylinux, _musllinux
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Collection, Iterable
+    from typing import AbstractSet
+
+
 __all__ = [
     "INTERPRETER_SHORT_NAMES",
     "AppleVersion",
@@ -50,6 +57,7 @@ def __dir__() -> list[str]:
 
 PythonVersion = Sequence[int]
 AppleVersion = Tuple[int, int]
+_T = TypeVar("_T")
 
 INTERPRETER_SHORT_NAMES: dict[str, str] = {
     "python": "py",  # Generic.
@@ -672,3 +680,43 @@ def sys_tags(*, warn: bool = False) -> Iterator[Tag]:
     else:
         interp = None
     yield from compatible_tags(interpreter=interp)
+
+
+def create_compatible_tags_selector(
+    tags: Iterable[Tag],
+) -> Callable[[Collection[tuple[_T, AbstractSet[Tag]]]], Sequence[_T]]:
+    """Create a callable to select things compatible with supported tags.
+
+    This function accepts an ordered sequence of tags, with the preferred
+    tags first.
+
+    The returned callable accepts a collection of tuples (thing, set[Tag]),
+    and returns an ordered sequence of things, with the things with the best
+    matching tags first. The returned sequence is empty if nothing matches.
+
+    Example to select compatible wheel filenames:
+
+    >>> filenames = ["foo-1.0-py3-none-any.whl", "foo-1.0-py2-none-any.whl"]
+    >>> selector = create_compatible_tags_selector(tags.sys_tags())
+    >>> compatible_filenames = selector([
+    ...     (filename, parse_wheel_filename(filename)[-1]) for filename in filenames
+    ... ])
+    ["foo-1.0-py3-none-any.whl"]
+    """
+    tag_ranks: dict[Tag, int] = {}
+    for rank, tag in enumerate(tags):
+        tag_ranks.setdefault(tag, rank)  # ignore duplicate tags, keep first
+    supported_tags = tag_ranks.keys()
+
+    def selector(
+        tagged_things: Collection[tuple[_T, AbstractSet[Tag]]],
+    ) -> Sequence[_T]:
+        ranked_things: list[tuple[_T, int]] = []
+        for thing, thing_tags in tagged_things:
+            supported_thing_tags = thing_tags & supported_tags
+            if supported_thing_tags:
+                thing_rank = min(tag_ranks[t] for t in supported_thing_tags)
+                ranked_things.append((thing, thing_rank))
+        return [thing for thing, _ in sorted(ranked_things, key=lambda rt: rt[1])]
+
+    return selector

tests/test_tags.py

@@ -1694,3 +1694,26 @@ def test_pickle() -> None:
     # Make sure equality works between a pickle/unpickle round trip.
     tag = tags.Tag("py3", "none", "any")
     assert pickle.loads(pickle.dumps(tag)) == tag
+
+
+@pytest.mark.parametrize(
+    ("supported", "things", "expected"),
+    [
+        (["t1", "t2"], ["t1", "t2"], ["t1", "t2"]),
+        (["t1", "t2"], ["t2", "t1"], ["t1", "t2"]),
+        (["t1", "t3"], ["t2", "t1"], ["t1"]),
+        (["t1", "t3"], ["t2.t3", "t1"], ["t1", "t2.t3"]),
+        (["t1"], ["t2", "t1"], ["t1"]),
+    ],
+)
+def test_create_compatible_tags_selector(
+    supported: list[str], things: list[str], expected: list[str]
+) -> None:
+    def t_to_tag(t: str) -> tags.Tag:
+        return tags.Tag("py3", "none", t)
+
+    def t_to_tags(t: str) -> frozenset[tags.Tag]:
+        return tags.parse_tag(f"py3-none-{t}")
+
+    selector = tags.create_compatible_tags_selector([t_to_tag(t) for t in supported])
+    assert selector([(t, t_to_tags(t)) for t in things]) == expected