xref: /aosp_15_r20/external/pytorch/tools/github/github_utils.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""GitHub Utilities"""
2
3from __future__ import annotations
4
5import json
6import os
7from typing import Any, Callable, cast, Dict
8from urllib.error import HTTPError
9from urllib.parse import quote
10from urllib.request import Request, urlopen
11
12
13def gh_fetch_url_and_headers(
14    url: str,
15    *,
16    headers: dict[str, str] | None = None,
17    data: dict[str, Any] | None = None,
18    method: str | None = None,
19    reader: Callable[[Any], Any] = lambda x: x.read(),
20) -> tuple[Any, Any]:
21    if headers is None:
22        headers = {}
23    token = os.environ.get("GITHUB_TOKEN")
24    if token is not None and url.startswith("https://api.github.com/"):
25        headers["Authorization"] = f"token {token}"
26    data_ = json.dumps(data).encode() if data is not None else None
27    try:
28        with urlopen(Request(url, headers=headers, data=data_, method=method)) as conn:
29            return conn.headers, reader(conn)
30    except HTTPError as err:
31        if err.code == 403 and all(
32            key in err.headers for key in ["X-RateLimit-Limit", "X-RateLimit-Used"]
33        ):
34            print(
35                f"""Rate limit exceeded:
36                Used: {err.headers['X-RateLimit-Used']}
37                Limit: {err.headers['X-RateLimit-Limit']}
38                Remaining: {err.headers['X-RateLimit-Remaining']}
39                Resets at: {err.headers['x-RateLimit-Reset']}"""
40            )
41        raise
42
43
44def gh_fetch_url(
45    url: str,
46    *,
47    headers: dict[str, str] | None = None,
48    data: dict[str, Any] | None = None,
49    method: str | None = None,
50    reader: Callable[[Any], Any] = lambda x: x.read(),
51) -> Any:
52    return gh_fetch_url_and_headers(
53        url, headers=headers, data=data, reader=json.load, method=method
54    )[1]
55
56
57def _gh_fetch_json_any(
58    url: str,
59    params: dict[str, Any] | None = None,
60    data: dict[str, Any] | None = None,
61) -> Any:
62    headers = {"Accept": "application/vnd.github.v3+json"}
63    if params is not None and len(params) > 0:
64        url += "?" + "&".join(
65            f"{name}={quote(str(val))}" for name, val in params.items()
66        )
67    return gh_fetch_url(url, headers=headers, data=data, reader=json.load)
68
69
70def gh_fetch_json_dict(
71    url: str,
72    params: dict[str, Any] | None = None,
73    data: dict[str, Any] | None = None,
74) -> dict[str, Any]:
75    return cast(Dict[str, Any], _gh_fetch_json_any(url, params, data))
76
77
78def gh_fetch_commit(org: str, repo: str, sha: str) -> dict[str, Any]:
79    return gh_fetch_json_dict(
80        f"https://api.github.com/repos/{org}/{repo}/commits/{sha}"
81    )
82