• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python3 -B
2# Copyright 2023 The Bazel Authors. All rights reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#    http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""A small script to update bazel files within the repo.
17
18We are not running this with 'bazel run' to keep the dependencies minimal
19"""
20
21# NOTE @aignas 2023-01-09: We should only depend on core Python 3 packages.
22import argparse
23import difflib
24import json
25import os
26import pathlib
27import sys
28import textwrap
29from collections import defaultdict
30from dataclasses import dataclass
31from typing import Any
32from urllib import request
33
34from tools.private.update_deps.args import path_from_runfiles
35from tools.private.update_deps.update_file import update_file
36
37# This should be kept in sync with //python:versions.bzl
38_supported_platforms = {
39    # Windows is unsupported right now
40    # "win_amd64": "x86_64-pc-windows-msvc",
41    "manylinux2014_x86_64": "x86_64-unknown-linux-gnu",
42    "manylinux2014_aarch64": "aarch64-unknown-linux-gnu",
43    "macosx_11_0_arm64": "aarch64-apple-darwin",
44    "macosx_10_9_x86_64": "x86_64-apple-darwin",
45    ("t", "manylinux2014_x86_64"): "x86_64-unknown-linux-gnu-freethreaded",
46    ("t", "manylinux2014_aarch64"): "aarch64-unknown-linux-gnu-freethreaded",
47    ("t", "macosx_11_0_arm64"): "aarch64-apple-darwin-freethreaded",
48    ("t", "macosx_10_9_x86_64"): "x86_64-apple-darwin-freethreaded",
49}
50
51
52@dataclass
53class Dep:
54    name: str
55    platform: str
56    python: str
57    url: str
58    sha256: str
59
60    @property
61    def repo_name(self):
62        return f"pypi__{self.name}_{self.python}_{self.platform}"
63
64    def __repr__(self):
65        return "\n".join(
66            [
67                "(",
68                f'    "{self.url}",',
69                f'    "{self.sha256}",',
70                ")",
71            ]
72        )
73
74
75@dataclass
76class Deps:
77    deps: list[Dep]
78
79    def __repr__(self):
80        deps = defaultdict(dict)
81        for d in self.deps:
82            deps[d.python][d.platform] = d
83
84        parts = []
85        for python, contents in deps.items():
86            inner = textwrap.indent(
87                "\n".join([f'"{platform}": {d},' for platform, d in contents.items()]),
88                prefix="    ",
89            )
90            parts.append('"{}": {{\n{}\n}},'.format(python, inner))
91        return "{{\n{}\n}}".format(textwrap.indent("\n".join(parts), prefix="    "))
92
93
94def _get_platforms(filename: str, python_version: str):
95    name, _, tail = filename.partition("-")
96    version, _, tail = tail.partition("-")
97    got_python_version, _, tail = tail.partition("-")
98    if python_version != got_python_version:
99        return []
100    abi, _, tail = tail.partition("-")
101
102    platforms, _, tail = tail.rpartition(".")
103    platforms = platforms.split(".")
104
105    return [("t", p) for p in platforms] if abi.endswith("t") else platforms
106
107
108def _map(
109    name: str,
110    filename: str,
111    python_version: str,
112    url: str,
113    digests: list,
114    platform: str,
115    **kwargs: Any,
116):
117    if platform not in _supported_platforms:
118        return None
119
120    return Dep(
121        name=name,
122        platform=_supported_platforms[platform],
123        python=python_version,
124        url=url,
125        sha256=digests["sha256"],
126    )
127
128
129def _parse_args() -> argparse.Namespace:
130    parser = argparse.ArgumentParser(__doc__)
131    parser.add_argument(
132        "--name",
133        default="coverage",
134        type=str,
135        help="The name of the package",
136    )
137    parser.add_argument(
138        "version",
139        type=str,
140        help="The version of the package to download",
141    )
142    parser.add_argument(
143        "--py",
144        nargs="+",
145        type=str,
146        default=["cp38", "cp39", "cp310", "cp311", "cp312", "cp313"],
147        help="Supported python versions",
148    )
149    parser.add_argument(
150        "--dry-run",
151        action="store_true",
152        help="Whether to write to files",
153    )
154    parser.add_argument(
155        "--update-file",
156        type=path_from_runfiles,
157        default=os.environ.get("UPDATE_FILE"),
158        help="The path for the file to be updated, defaults to the value taken from UPDATE_FILE",
159    )
160    return parser.parse_args()
161
162
163def main():
164    args = _parse_args()
165
166    api_url = f"https://pypi.org/pypi/{args.name}/{args.version}/json"
167    req = request.Request(api_url)
168    with request.urlopen(req) as response:
169        data = json.loads(response.read().decode("utf-8"))
170
171    urls = []
172    for u in data["urls"]:
173        if u["yanked"]:
174            continue
175
176        if not u["filename"].endswith(".whl"):
177            continue
178
179        if u["python_version"] not in args.py:
180            continue
181
182        if f'_{u["python_version"]}m_' in u["filename"]:
183            continue
184
185        platforms = _get_platforms(
186            u["filename"],
187            u["python_version"],
188        )
189
190        result = [_map(name=args.name, platform=p, **u) for p in platforms]
191        urls.extend(filter(None, result))
192
193    urls.sort(key=lambda x: f"{x.python}_{x.platform}")
194
195    # Update the coverage_deps, which are used to register deps
196    update_file(
197        path=args.update_file,
198        snippet=f"_coverage_deps = {repr(Deps(urls))}\n",
199        start_marker="# START: maintained by 'bazel run //tools/private/update_deps:update_coverage_deps <version>'",
200        end_marker="# END: maintained by 'bazel run //tools/private/update_deps:update_coverage_deps <version>'",
201        dry_run=args.dry_run,
202    )
203
204    return
205
206
207if __name__ == "__main__":
208    main()
209