pytorch

1
"""GitHub Utilities"""
2

3
import json
4
import os
5
import warnings
6
from dataclasses import dataclass
7
from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
8
from urllib.error import HTTPError
9
from urllib.parse import quote
10
from urllib.request import Request, urlopen
11

12

13
GITHUB_API_URL = "https://api.github.com"
14

15

16
@dataclass
17
class GitHubComment:
18
    body_text: str
19
    created_at: str
20
    author_login: str
21
    author_association: str
22
    editor_login: Optional[str]
23
    database_id: int
24
    url: str
25

26

27
def gh_fetch_url_and_headers(
28
    url: str,
29
    *,
30
    headers: Optional[Dict[str, str]] = None,
31
    data: Union[Optional[Dict[str, Any]], str] = None,
32
    method: Optional[str] = None,
33
    reader: Callable[[Any], Any] = lambda x: x.read(),
34
) -> Tuple[Any, Any]:
35
    if headers is None:
36
        headers = {}
37
    token = os.environ.get("GITHUB_TOKEN")
38
    if token is not None and url.startswith(f"{GITHUB_API_URL}/"):
39
        headers["Authorization"] = f"token {token}"
40

41
    data_ = None
42
    if data is not None:
43
        data_ = data.encode() if isinstance(data, str) else json.dumps(data).encode()
44

45
    try:
46
        with urlopen(Request(url, headers=headers, data=data_, method=method)) as conn:
47
            return conn.headers, reader(conn)
48
    except HTTPError as err:
49
        if (
50
            err.code == 403
51
            and all(
52
                key in err.headers
53
                for key in ["X-RateLimit-Limit", "X-RateLimit-Remaining"]
54
            )
55
            and int(err.headers["X-RateLimit-Remaining"]) == 0
56
        ):
57
            print(
58
                f"""{url}
59
                Rate limit exceeded:
60
                Used: {err.headers['X-RateLimit-Used']}
61
                Limit: {err.headers['X-RateLimit-Limit']}
62
                Remaining: {err.headers['X-RateLimit-Remaining']}
63
                Resets at: {err.headers['x-RateLimit-Reset']}"""
64
            )
65
        else:
66
            print(f"Error fetching {url} {err}")
67
        raise
68

69

70
def gh_fetch_url(
71
    url: str,
72
    *,
73
    headers: Optional[Dict[str, str]] = None,
74
    data: Union[Optional[Dict[str, Any]], str] = None,
75
    method: Optional[str] = None,
76
    reader: Callable[[Any], Any] = lambda x: x.read(),
77
) -> Any:
78
    return gh_fetch_url_and_headers(
79
        url, headers=headers, data=data, reader=json.load, method=method
80
    )[1]
81

82

83
def gh_fetch_json(
84
    url: str,
85
    params: Optional[Dict[str, Any]] = None,
86
    data: Optional[Dict[str, Any]] = None,
87
    method: Optional[str] = None,
88
) -> List[Dict[str, Any]]:
89
    headers = {"Accept": "application/vnd.github.v3+json"}
90
    if params is not None and len(params) > 0:
91
        url += "?" + "&".join(
92
            f"{name}={quote(str(val))}" for name, val in params.items()
93
        )
94
    return cast(
95
        List[Dict[str, Any]],
96
        gh_fetch_url(url, headers=headers, data=data, reader=json.load, method=method),
97
    )
98

99

100
def _gh_fetch_json_any(
101
    url: str,
102
    params: Optional[Dict[str, Any]] = None,
103
    data: Optional[Dict[str, Any]] = None,
104
) -> Any:
105
    headers = {"Accept": "application/vnd.github.v3+json"}
106
    if params is not None and len(params) > 0:
107
        url += "?" + "&".join(
108
            f"{name}={quote(str(val))}" for name, val in params.items()
109
        )
110
    return gh_fetch_url(url, headers=headers, data=data, reader=json.load)
111

112

113
def gh_fetch_json_list(
114
    url: str,
115
    params: Optional[Dict[str, Any]] = None,
116
    data: Optional[Dict[str, Any]] = None,
117
) -> List[Dict[str, Any]]:
118
    return cast(List[Dict[str, Any]], _gh_fetch_json_any(url, params, data))
119

120

121
def gh_fetch_json_dict(
122
    url: str,
123
    params: Optional[Dict[str, Any]] = None,
124
    data: Optional[Dict[str, Any]] = None,
125
) -> Dict[str, Any]:
126
    return cast(Dict[str, Any], _gh_fetch_json_any(url, params, data))
127

128

129
def gh_graphql(query: str, **kwargs: Any) -> Dict[str, Any]:
130
    rc = gh_fetch_url(
131
        "https://api.github.com/graphql",
132
        data={"query": query, "variables": kwargs},
133
        reader=json.load,
134
    )
135
    if "errors" in rc:
136
        raise RuntimeError(
137
            f"GraphQL query {query}, args {kwargs} failed: {rc['errors']}"
138
        )
139
    return cast(Dict[str, Any], rc)
140

141

142
def _gh_post_comment(
143
    url: str, comment: str, dry_run: bool = False
144
) -> List[Dict[str, Any]]:
145
    if dry_run:
146
        print(comment)
147
        return []
148
    return gh_fetch_json_list(url, data={"body": comment})
149

150

151
def gh_post_pr_comment(
152
    org: str, repo: str, pr_num: int, comment: str, dry_run: bool = False
153
) -> List[Dict[str, Any]]:
154
    return _gh_post_comment(
155
        f"{GITHUB_API_URL}/repos/{org}/{repo}/issues/{pr_num}/comments",
156
        comment,
157
        dry_run,
158
    )
159

160

161
def gh_post_commit_comment(
162
    org: str, repo: str, sha: str, comment: str, dry_run: bool = False
163
) -> List[Dict[str, Any]]:
164
    return _gh_post_comment(
165
        f"{GITHUB_API_URL}/repos/{org}/{repo}/commits/{sha}/comments",
166
        comment,
167
        dry_run,
168
    )
169

170

171
def gh_delete_comment(org: str, repo: str, comment_id: int) -> None:
172
    url = f"{GITHUB_API_URL}/repos/{org}/{repo}/issues/comments/{comment_id}"
173
    gh_fetch_url(url, method="DELETE")
174

175

176
def gh_fetch_merge_base(org: str, repo: str, base: str, head: str) -> str:
177
    merge_base = ""
178
    # Get the merge base using the GitHub REST API. This is the same as using
179
    # git merge-base without the need to have git. The API doc can be found at
180
    # https://docs.github.com/en/rest/commits/commits?apiVersion=2022-11-28#compare-two-commits
181
    try:
182
        json_data = gh_fetch_url(
183
            f"{GITHUB_API_URL}/repos/{org}/{repo}/compare/{base}...{head}",
184
            headers={"Accept": "application/vnd.github.v3+json"},
185
            reader=json.load,
186
        )
187
        if json_data:
188
            merge_base = json_data.get("merge_base_commit", {}).get("sha", "")
189
        else:
190
            warnings.warn(
191
                f"Failed to get merge base for {base}...{head}: Empty response"
192
            )
193
    except Exception as error:
194
        warnings.warn(f"Failed to get merge base for {base}...{head}: {error}")
195

196
    return merge_base
197

198

199
def gh_update_pr_state(org: str, repo: str, pr_num: int, state: str = "open") -> None:
200
    url = f"{GITHUB_API_URL}/repos/{org}/{repo}/pulls/{pr_num}"
201
    try:
202
        gh_fetch_url(url, method="PATCH", data={"state": state})
203
    except HTTPError as err:
204
        # When trying to open the pull request, error 422 means that the branch
205
        # has been deleted and the API couldn't re-open it
206
        if err.code == 422 and state == "open":
207
            warnings.warn(
208
                f"Failed to open {pr_num} because its head branch has been deleted: {err}"
209
            )
210
        else:
211
            raise
212

213

214
def gh_query_issues_by_labels(
215
    org: str, repo: str, labels: List[str], state: str = "open"
216
) -> List[Dict[str, Any]]:
217
    url = f"{GITHUB_API_URL}/repos/{org}/{repo}/issues"
218
    return gh_fetch_json(
219
        url, method="GET", params={"labels": ",".join(labels), "state": state}
220
    )
221