refactor/cleanup code

Author: aignas

python/private/pypi/select_whl.bzl

@@ -4,55 +4,80 @@ load("//python/private:version.bzl", "version")
 load(":parse_whl_name.bzl", "parse_whl_name")
 load(":python_tag.bzl", "PY_TAG_GENERIC", "python_tag")
 
-def _get_priority(*, tags, values, allow_wildcard = True):
+def _get_priority(*, tag, values, allow_wildcard = True):
     keys = []
     for priority, wp in enumerate(values):
-        for tag in tags.split("."):
-            head, sep, tail = wp.partition("*")
-            if "*" in tail:
-                fail("only a single '*' can be present in the matcher")
-            if not allow_wildcard and sep:
-                fail("'*' is not allowed in the matcher")
-
-            if not sep and tag == head:
-                keys.append(priority)
-            elif sep and tag.startswith(head) and tag.endswith(tail):
-                keys.append(priority)
+        head, sep, tail = wp.partition("*")
+        if "*" in tail:
+            fail("only a single '*' can be present in the matcher")
+        if not allow_wildcard and sep:
+            fail("'*' is not allowed in the matcher")
+
+        if not sep and tag == head:
+            keys.append(priority)
+        elif sep and tag.startswith(head) and tag.endswith(tail):
+            keys.append(priority)
 
     if not keys:
         return None
 
     return max(keys)
 
-def _get_py_priority(*, tags, implementation, py_version):
-    keys = []
-    for tag in tags.split("."):
-        if tag.startswith(PY_TAG_GENERIC):
-            ver_str = tag[len(PY_TAG_GENERIC):]
-        elif tag.startswith(implementation):
-            ver_str = tag[len(implementation):]
-        else:
-            continue
+def _get_py_priority(*, tag, implementation, py_version):
+    if tag.startswith(PY_TAG_GENERIC):
+        ver_str = tag[len(PY_TAG_GENERIC):]
+    elif tag.startswith(implementation):
+        ver_str = tag[len(implementation):]
+    else:
+        return None
 
-        # Add a 0 at the end in case it is a single digit
-        ver_str = "{}.{}".format(ver_str[0], ver_str[1:] or "0")
+    # Add a 0 at the end in case it is a single digit
+    ver_str = "{}.{}".format(ver_str[0], ver_str[1:] or "0")
 
-        ver = version.parse(ver_str)
-        if not version.is_compatible(py_version, ver):
-            continue
+    ver = version.parse(ver_str)
+    if not version.is_compatible(py_version, ver):
+        return None
 
-        keys.append((
-            tag.startswith(implementation),
-            version.key(ver),
-            # Prefer shorter py_tags, which will yield more specialized matches,
-            # like preferring py3 over py2.py3
-            -len(tags),
-        ))
+    return (
+        tag.startswith(implementation),
+        version.key(ver),
+    )
 
-    if not keys:
-        return None
+def _tag_sets(*, whls, implementation, py_version, whl_abi_tags, platforms):
+    ret = {}
+    for whl in whls:
+        parsed = parse_whl_name(whl.filename)
 
-    return max(keys)
+        # See https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/#compressed-tag-sets
+        for py in parsed.python_tag.split("."):
+            py = _get_py_priority(
+                tag = py,
+                implementation = implementation,
+                py_version = py_version,
+            )
+            if py == None:
+                ret[struct(py = py)] = whl
+                continue
+
+            for abi in parsed.abi_tag.split("."):
+                abi = _get_priority(
+                    tag = abi,
+                    values = whl_abi_tags,
+                    allow_wildcard = False,
+                )
+                if py == None:
+                    ret[struct(py = py, abi = abi)] = whl
+                    continue
+
+                for p in parsed.platform_tag.split("."):
+                    platform = _get_priority(
+                        tag = p,
+                        values = platforms,
+                    )
+
+                    ret[struct(py = py, abi = abi, platform = platform)] = whl
+
+    return ret
 
 def select_whl(*, whls, python_version, platforms, whl_abi_tags, implementation_name = "cpython", limit = 1, logger = None):
     """Select a whl that is the most suitable for the given platform.
@@ -77,65 +102,41 @@ def select_whl(*, whls, python_version, platforms, whl_abi_tags, implementation_
     candidates = {}
     implementation = python_tag(implementation_name)
 
-    for whl in whls:
-        parsed = parse_whl_name(whl.filename)
-
-        if parsed.python_tag.startswith(PY_TAG_GENERIC):
-            pass
-        elif not parsed.python_tag.startswith(implementation):
+    for priority, whl in _tag_sets(
+        whls = whls,
+        implementation = implementation,
+        py_version = py_version,
+        whl_abi_tags = whl_abi_tags,
+        platforms = platforms,
+    ).items():
+        if priority.py == None:
             if logger:
-                logger.debug(lambda: "Discarding the wheel because the implementation '{}' is not compatible with target implementation '{}'".format(
-                    parsed.python_tag,
-                    implementation,
-                ))
-            continue
-
-        py_priority = _get_py_priority(
-            tags = parsed.python_tag,
-            implementation = implementation,
-            py_version = py_version,
-        )
-        if py_priority == None:
-            if logger:
-                logger.debug(lambda: "The py_tag '{}' does not match implementation version: {} {}".format(
-                    parsed.py_tag,
+                logger.debug(lambda: "The python_tag in '{}' does not match implementation or version: {} {}".format(
+                    whl.filename,
                     implementation,
                     py_version.string,
                 ))
             continue
-
-        abi_priority = _get_priority(
-            tags = parsed.abi_tag,
-            values = whl_abi_tags,
-            allow_wildcard = False,
-        )
-        if abi_priority == None:
+        elif priority.abi == None:
             if logger:
-                logger.debug(lambda: "The abi '{}' does not match given list: {}".format(
-                    parsed.abi_tag,
+                logger.debug(lambda: "The abi_tag in '{}' does not match given list: {}".format(
+                    whl.filename,
                     whl_abi_tags,
                 ))
             continue
-
-        platform_priority = _get_priority(
-            tags = parsed.platform_tag,
-            values = platforms,
-        )
-        if platform_priority == None:
+        elif priority.platform == None:
             if logger:
-                logger.debug(lambda: "The platform_tag '{}' does not match given list: {}".format(
-                    parsed.platform_tag,
+                logger.debug(lambda: "The platform_tag in '{}' does not match given list: {}".format(
+                    whl.filename,
                     platforms,
                 ))
             continue
 
-        key = (
-            # Ensure that we chose the highest compatible version
-            py_priority,
-            platform_priority,
-            abi_priority,
-        )
-        candidates.setdefault(key, whl)
+        candidates.setdefault((
+            priority.platform,  # Prefer platform wheels
+            priority.py,  # Then prefer implementation/python version
+            priority.abi,  # Then prefer more specific ABI wheels
+        ), whl)
 
     if not candidates:
         return None