xref: /aosp_15_r20/external/bazelbuild-rules_python/tools/private/update_deps/update_coverage_deps.py (revision 60517a1edbc8ecf509223e9af94a7adec7d736b8) 
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}
46
47
48@dataclass
49class Dep:
50    name: str
51    platform: str
52    python: str
53    url: str
54    sha256: str
55
56    @property
57    def repo_name(self):
58        return f"pypi__{self.name}_{self.python}_{self.platform}"
59
60    def __repr__(self):
61        return "\n".join(
62            [
63                "(",
64                f'    "{self.url}",',
65                f'    "{self.sha256}",',
66                ")",
67            ]
68        )
69
70
71@dataclass
72class Deps:
73    deps: list[Dep]
74
75    def __repr__(self):
76        deps = defaultdict(dict)
77        for d in self.deps:
78            deps[d.python][d.platform] = d
79
80        parts = []
81        for python, contents in deps.items():
82            inner = textwrap.indent(
83                "\n".join([f'"{platform}": {d},' for platform, d in contents.items()]),
84                prefix="    ",
85            )
86            parts.append('"{}": {{\n{}\n}},'.format(python, inner))
87        return "{{\n{}\n}}".format(textwrap.indent("\n".join(parts), prefix="    "))
88
89
90def _get_platforms(filename: str, name: str, version: str, python_version: str):
91    return filename[
92        len(f"{name}-{version}-{python_version}-{python_version}-") : -len(".whl")
93    ].split(".")
94
95
96def _map(
97    name: str,
98    filename: str,
99    python_version: str,
100    url: str,
101    digests: list,
102    platform: str,
103    **kwargs: Any,
104):
105    if platform not in _supported_platforms:
106        return None
107
108    return Dep(
109        name=name,
110        platform=_supported_platforms[platform],
111        python=python_version,
112        url=url,
113        sha256=digests["sha256"],
114    )
115
116
117def _parse_args() -> argparse.Namespace:
118    parser = argparse.ArgumentParser(__doc__)
119    parser.add_argument(
120        "--name",
121        default="coverage",
122        type=str,
123        help="The name of the package",
124    )
125    parser.add_argument(
126        "version",
127        type=str,
128        help="The version of the package to download",
129    )
130    parser.add_argument(
131        "--py",
132        nargs="+",
133        type=str,
134        default=["cp38", "cp39", "cp310", "cp311", "cp312"],
135        help="Supported python versions",
136    )
137    parser.add_argument(
138        "--dry-run",
139        action="store_true",
140        help="Whether to write to files",
141    )
142    parser.add_argument(
143        "--update-file",
144        type=path_from_runfiles,
145        default=os.environ.get("UPDATE_FILE"),
146        help="The path for the file to be updated, defaults to the value taken from UPDATE_FILE",
147    )
148    return parser.parse_args()
149
150
151def main():
152    args = _parse_args()
153
154    api_url = f"https://pypi.org/pypi/{args.name}/{args.version}/json"
155    req = request.Request(api_url)
156    with request.urlopen(req) as response:
157        data = json.loads(response.read().decode("utf-8"))
158
159    urls = []
160    for u in data["urls"]:
161        if u["yanked"]:
162            continue
163
164        if not u["filename"].endswith(".whl"):
165            continue
166
167        if u["python_version"] not in args.py:
168            continue
169
170        if f'_{u["python_version"]}m_' in u["filename"]:
171            continue
172
173        platforms = _get_platforms(
174            u["filename"],
175            args.name,
176            args.version,
177            u["python_version"],
178        )
179
180        result = [_map(name=args.name, platform=p, **u) for p in platforms]
181        urls.extend(filter(None, result))
182
183    urls.sort(key=lambda x: f"{x.python}_{x.platform}")
184
185    # Update the coverage_deps, which are used to register deps
186    update_file(
187        path=args.update_file,
188        snippet=f"_coverage_deps = {repr(Deps(urls))}\n",
189        start_marker="# START: maintained by 'bazel run //tools/private/update_deps:update_coverage_deps <version>'",
190        end_marker="# END: maintained by 'bazel run //tools/private/update_deps:update_coverage_deps <version>'",
191        dry_run=args.dry_run,
192    )
193
194    return
195
196
197if __name__ == "__main__":
198    main()
199