• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2023 The Bazel Authors. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""
16A starlark implementation of the wheel platform tag parsing to get the target platform.
17"""
18
19load(":parse_whl_name.bzl", "parse_whl_name")
20
21# The order of the dictionaries is to keep definitions with their aliases next to each
22# other
23_CPU_ALIASES = {
24    "x86_32": "x86_32",
25    "i386": "x86_32",
26    "i686": "x86_32",
27    "x86": "x86_32",
28    "x86_64": "x86_64",
29    "amd64": "x86_64",
30    "aarch64": "aarch64",
31    "arm64": "aarch64",
32    "ppc": "ppc",
33    "ppc64": "ppc",
34    "ppc64le": "ppc",
35    "s390x": "s390x",
36    "arm": "arm",
37    "armv6l": "arm",
38    "armv7l": "arm",
39}  # buildifier: disable=unsorted-dict-items
40
41_OS_PREFIXES = {
42    "linux": "linux",
43    "manylinux": "linux",
44    "musllinux": "linux",
45    "macos": "osx",
46    "win": "windows",
47}  # buildifier: disable=unsorted-dict-items
48
49def select_whls(*, whls, want_platforms = [], logger = None):
50    """Select a subset of wheels suitable for target platforms from a list.
51
52    Args:
53        whls(list[struct]): A list of candidates which have a `filename`
54            attribute containing the `whl` filename.
55        want_platforms(str): The platforms in "{abi}_{os}_{cpu}" or "{os}_{cpu}" format.
56        logger: A logger for printing diagnostic messages.
57
58    Returns:
59        A filtered list of items from the `whls` arg where `filename` matches
60        the selected criteria. If no match is found, an empty list is returned.
61    """
62    if not whls:
63        return []
64
65    want_abis = {
66        "abi3": None,
67        "none": None,
68    }
69
70    _want_platforms = {}
71    version_limit = None
72
73    for p in want_platforms:
74        if not p.startswith("cp3"):
75            fail("expected all platforms to start with ABI, but got: {}".format(p))
76
77        abi, _, os_cpu = p.partition("_")
78        _want_platforms[os_cpu] = None
79        _want_platforms[p] = None
80
81        version_limit_candidate = int(abi[3:])
82        if not version_limit:
83            version_limit = version_limit_candidate
84        if version_limit and version_limit != version_limit_candidate:
85            fail("Only a single python version is supported for now")
86
87        # For some legacy implementations the wheels may target the `cp3xm` ABI
88        _want_platforms["{}m_{}".format(abi, os_cpu)] = None
89        want_abis[abi] = None
90        want_abis[abi + "m"] = None
91
92    want_platforms = sorted(_want_platforms)
93
94    candidates = {}
95    for whl in whls:
96        parsed = parse_whl_name(whl.filename)
97
98        if logger:
99            logger.trace(lambda: "Deciding whether to use '{}'".format(whl.filename))
100
101        supported_implementations = {}
102        whl_version_min = 0
103        for tag in parsed.python_tag.split("."):
104            supported_implementations[tag[:2]] = None
105
106            if tag.startswith("cp3") or tag.startswith("py3"):
107                version = int(tag[len("..3"):] or 0)
108            else:
109                # In this case it should be eithor "cp2" or "py2" and we will default
110                # to `whl_version_min` = 0
111                continue
112
113            if whl_version_min == 0 or version < whl_version_min:
114                whl_version_min = version
115
116        if not ("cp" in supported_implementations or "py" in supported_implementations):
117            if logger:
118                logger.trace(lambda: "Discarding the whl because the whl does not support CPython, whl supported implementations are: {}".format(supported_implementations))
119            continue
120
121        if want_abis and parsed.abi_tag not in want_abis:
122            # Filter out incompatible ABIs
123            if logger:
124                logger.trace(lambda: "Discarding the whl because the whl abi did not match")
125            continue
126
127        if whl_version_min > version_limit:
128            if logger:
129                logger.trace(lambda: "Discarding the whl because the whl supported python version is too high")
130            continue
131
132        compatible = False
133        if parsed.platform_tag == "any":
134            compatible = True
135        else:
136            for p in whl_target_platforms(parsed.platform_tag, abi_tag = parsed.abi_tag.strip("m") if parsed.abi_tag.startswith("cp") else None):
137                if p.target_platform in want_platforms:
138                    compatible = True
139                    break
140
141        if not compatible:
142            if logger:
143                logger.trace(lambda: "Discarding the whl because the whl does not support the desired platforms: {}".format(want_platforms))
144            continue
145
146        for implementation in supported_implementations:
147            candidates.setdefault(
148                (
149                    parsed.abi_tag,
150                    parsed.platform_tag,
151                ),
152                {},
153            ).setdefault(
154                (
155                    # prefer cp implementation
156                    implementation == "cp",
157                    # prefer higher versions
158                    whl_version_min,
159                    # prefer abi3 over none
160                    parsed.abi_tag != "none",
161                    # prefer cpx abi over abi3
162                    parsed.abi_tag != "abi3",
163                ),
164                [],
165            ).append(whl)
166
167    return [
168        candidates[key][sorted(v)[-1]][-1]
169        for key, v in candidates.items()
170    ]
171
172def whl_target_platforms(platform_tag, abi_tag = ""):
173    """Parse the wheel abi and platform tags and return (os, cpu) tuples.
174
175    Args:
176        platform_tag (str): The platform_tag part of the wheel name. See
177            ./parse_whl_name.bzl for more details.
178        abi_tag (str): The abi tag that should be used for parsing.
179
180    Returns:
181        A list of structs, with attributes:
182        * os: str, one of the _OS_PREFIXES values
183        * cpu: str, one of the _CPU_PREFIXES values
184        * abi: str, the ABI that the interpreter should have if it is passed.
185        * target_platform: str, the target_platform that can be given to the
186          wheel_installer for parsing whl METADATA.
187    """
188    cpus = _cpu_from_tag(platform_tag)
189
190    abi = None
191    if abi_tag not in ["", "none", "abi3"]:
192        abi = abi_tag
193
194    # TODO @aignas 2024-05-29: this code is present in many places, I think
195    _, _, tail = platform_tag.partition("_")
196    maybe_arch = tail
197    major, _, tail = tail.partition("_")
198    minor, _, tail = tail.partition("_")
199    if not tail or not major.isdigit() or not minor.isdigit():
200        tail = maybe_arch
201        major = 0
202        minor = 0
203
204    for prefix, os in _OS_PREFIXES.items():
205        if platform_tag.startswith(prefix):
206            return [
207                struct(
208                    os = os,
209                    cpu = cpu,
210                    abi = abi,
211                    version = (int(major), int(minor)),
212                    target_platform = "_".join([abi, os, cpu] if abi else [os, cpu]),
213                )
214                for cpu in cpus
215            ]
216
217    print("WARNING: ignoring unknown platform_tag os: {}".format(platform_tag))  # buildifier: disable=print
218    return []
219
220def _cpu_from_tag(tag):
221    candidate = [
222        cpu
223        for input, cpu in _CPU_ALIASES.items()
224        if tag.endswith(input)
225    ]
226    if candidate:
227        return candidate
228
229    if tag == "win32":
230        return ["x86_32"]
231    elif tag == "win_ia64":
232        return []
233    elif tag.startswith("macosx"):
234        if tag.endswith("universal2"):
235            return ["x86_64", "aarch64"]
236        elif tag.endswith("universal"):
237            return ["x86_64", "aarch64"]
238        elif tag.endswith("intel"):
239            return ["x86_32"]
240
241    return []
242