|
| 1 | +# /// script |
| 2 | +# requires-python = ">=3.13" |
| 3 | +# dependencies = [] |
| 4 | +# /// |
| 5 | +"""Request a reviewer from repository pools using shared round-robin credits.""" |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import argparse |
| 10 | +import json |
| 11 | +import math |
| 12 | +import os |
| 13 | +import random |
| 14 | +import sys |
| 15 | +import tomllib |
| 16 | +from dataclasses import dataclass |
| 17 | +from pathlib import Path, PurePosixPath |
| 18 | +from typing import Any |
| 19 | +from urllib.error import HTTPError |
| 20 | +from urllib.request import Request, urlopen |
| 21 | + |
| 22 | +REPOSITORY = "astral-sh/ruff" |
| 23 | +REPOSITORY_ID = 523043277 |
| 24 | +API_ROOT = f"https://api.github.com/repos/{REPOSITORY}" |
| 25 | +State = dict[str, dict[str, float]] |
| 26 | + |
| 27 | + |
| 28 | +@dataclass(frozen=True) |
| 29 | +class Pool: |
| 30 | + name: str |
| 31 | + patterns: tuple[str, ...] |
| 32 | + reviewers: tuple[str, ...] |
| 33 | + |
| 34 | + |
| 35 | +def api(endpoint: str, *, data: dict | None = None) -> Any: |
| 36 | + request = Request( |
| 37 | + f"{API_ROOT}/{endpoint}", |
| 38 | + data=json.dumps(data).encode() if data is not None else None, |
| 39 | + headers={ |
| 40 | + "Authorization": f"Bearer {os.environ['GH_TOKEN']}", |
| 41 | + "Accept": "application/vnd.github+json", |
| 42 | + "Content-Type": "application/json", |
| 43 | + "X-GitHub-Api-Version": "2022-11-28", |
| 44 | + }, |
| 45 | + method="POST" if data is not None else "GET", |
| 46 | + ) |
| 47 | + with urlopen(request, timeout=30) as response: |
| 48 | + return json.load(response) |
| 49 | + |
| 50 | + |
| 51 | +def changed_files(pull_request: int) -> list[str]: |
| 52 | + files = [] |
| 53 | + page = 1 |
| 54 | + while True: |
| 55 | + batch = api(f"pulls/{pull_request}/files?per_page=100&page={page}") |
| 56 | + files.extend(item["filename"] for item in batch) |
| 57 | + if len(batch) < 100: |
| 58 | + return files |
| 59 | + page += 1 |
| 60 | + |
| 61 | + |
| 62 | +def parse_pools(contents: str) -> list[Pool]: |
| 63 | + return [ |
| 64 | + Pool( |
| 65 | + name=pool["name"], |
| 66 | + patterns=tuple(path.lstrip("/") for path in pool["paths"]), |
| 67 | + reviewers=tuple(pool["reviewers"]), |
| 68 | + ) |
| 69 | + for pool in tomllib.loads(contents)["pools"] |
| 70 | + ] |
| 71 | + |
| 72 | + |
| 73 | +def matching_pools(pools: list[Pool], files: list[str]) -> list[tuple[Pool, int]]: |
| 74 | + matched = [] |
| 75 | + for pool in pools: |
| 76 | + weight = sum( |
| 77 | + any( |
| 78 | + PurePosixPath(filename).full_match(pattern) for pattern in pool.patterns |
| 79 | + ) |
| 80 | + for filename in files |
| 81 | + ) |
| 82 | + if weight: |
| 83 | + matched.append((pool, weight)) |
| 84 | + return matched |
| 85 | + |
| 86 | + |
| 87 | +def parse_state(contents: str) -> State: |
| 88 | + parsed = json.loads(contents) |
| 89 | + if not isinstance(parsed, dict): |
| 90 | + raise TypeError("Reviewer state must be an object.") |
| 91 | + state: State = {} |
| 92 | + for name, value in parsed.items(): |
| 93 | + if not isinstance(value, dict): |
| 94 | + raise TypeError(f"Reviewer state for pool {name} must be an object.") |
| 95 | + credits = {} |
| 96 | + for login, credit in value.items(): |
| 97 | + if type(credit) not in (int, float) or not math.isfinite(credit): |
| 98 | + raise ValueError(f"Invalid reviewer credit in pool {name}.") |
| 99 | + credits[login] = float(credit) |
| 100 | + state[name] = credits |
| 101 | + return state |
| 102 | + |
| 103 | + |
| 104 | +def load_state(path: Path) -> State: |
| 105 | + if path.exists(): |
| 106 | + return parse_state(path.read_text()) |
| 107 | + return {} |
| 108 | + |
| 109 | + |
| 110 | +def eligible_pull(pull: dict, pull_request: int, head_sha: str) -> bool: |
| 111 | + return ( |
| 112 | + pull["number"] == pull_request |
| 113 | + and pull["state"] == "open" |
| 114 | + and not pull["draft"] |
| 115 | + and pull["base"]["repo"]["id"] == REPOSITORY_ID |
| 116 | + and (pull["head"]["repo"] or {}).get("id") == REPOSITORY_ID |
| 117 | + and pull["head"]["sha"] == head_sha |
| 118 | + ) |
| 119 | + |
| 120 | + |
| 121 | +def has_reviewer(pull: dict) -> bool: |
| 122 | + author = pull["user"]["login"].lower() |
| 123 | + return any(user["login"].lower() != author for user in pull["requested_reviewers"]) |
| 124 | + |
| 125 | + |
| 126 | +def assign_reviewer( |
| 127 | + pull_request: int, |
| 128 | + head_sha: str, |
| 129 | + pools: list[Pool], |
| 130 | + state: State, |
| 131 | + *, |
| 132 | + dry_run: bool = False, |
| 133 | + rng: random.Random | None = None, |
| 134 | +) -> tuple[str, str] | None: |
| 135 | + rng = rng or random.Random() |
| 136 | + pull = api(f"pulls/{pull_request}") |
| 137 | + if not eligible_pull(pull, pull_request, head_sha) or has_reviewer(pull): |
| 138 | + return None |
| 139 | + author = pull["user"]["login"].lower() |
| 140 | + matched = matching_pools(pools, changed_files(pull_request)) |
| 141 | + while matched: |
| 142 | + eligible_pools, weights = zip(*matched, strict=True) |
| 143 | + pool = rng.choices(eligible_pools, weights=weights)[0] |
| 144 | + stored = state.get(pool.name, {}) |
| 145 | + credits = { |
| 146 | + login: min(1.0, stored.get(login, 0.0) + 1 / len(pool.reviewers)) |
| 147 | + for login in pool.reviewers |
| 148 | + } |
| 149 | + candidates = [login for login in pool.reviewers if login.lower() != author] |
| 150 | + while candidates: |
| 151 | + highest = max(credits[login] for login in candidates) |
| 152 | + reviewer = rng.choice( |
| 153 | + [login for login in candidates if credits[login] == highest] |
| 154 | + ) |
| 155 | + candidates.remove(reviewer) |
| 156 | + current = api(f"pulls/{pull_request}") |
| 157 | + if not eligible_pull(current, pull_request, head_sha) or has_reviewer( |
| 158 | + current |
| 159 | + ): |
| 160 | + return None |
| 161 | + if dry_run: |
| 162 | + print(f"Would request review from {reviewer} in pool {pool.name}.") |
| 163 | + return pool.name, reviewer |
| 164 | + try: |
| 165 | + api( |
| 166 | + f"pulls/{pull_request}/requested_reviewers", |
| 167 | + data={"reviewers": [reviewer]}, |
| 168 | + ) |
| 169 | + except HTTPError as error: |
| 170 | + error.close() |
| 171 | + if error.code == 422: |
| 172 | + print(f"GitHub rejected reviewer {reviewer}.", file=sys.stderr) |
| 173 | + continue |
| 174 | + raise |
| 175 | + credits[reviewer] = max(-1.0, credits[reviewer] - 1) |
| 176 | + state[pool.name] = credits |
| 177 | + print(f"Requested review from {reviewer} in pool {pool.name}.") |
| 178 | + return pool.name, reviewer |
| 179 | + matched = [entry for entry in matched if entry[0] != pool] |
| 180 | + return None |
| 181 | + |
| 182 | + |
| 183 | +def main() -> None: |
| 184 | + parser = argparse.ArgumentParser(description=__doc__) |
| 185 | + parser.add_argument("--pull-request", required=True, type=int) |
| 186 | + parser.add_argument("--head-sha", required=True) |
| 187 | + parser.add_argument("--state", required=True, type=Path) |
| 188 | + parser.add_argument("--output", type=Path) |
| 189 | + parser.add_argument("--dry-run", action="store_true") |
| 190 | + args = parser.parse_args() |
| 191 | + if os.environ.get("GITHUB_REPOSITORY", REPOSITORY) != REPOSITORY: |
| 192 | + parser.error(f"This script only supports {REPOSITORY}.") |
| 193 | + if args.pull_request <= 0: |
| 194 | + parser.error("The pull request number must be positive.") |
| 195 | + if not args.dry_run and args.output is None: |
| 196 | + parser.error("--output is required when requesting a review.") |
| 197 | + pools = parse_pools(Path(".github/pr-reviewer-pools.toml").read_text()) |
| 198 | + state = load_state(args.state) |
| 199 | + assign_reviewer( |
| 200 | + args.pull_request, args.head_sha, pools, state, dry_run=args.dry_run |
| 201 | + ) |
| 202 | + if not args.dry_run: |
| 203 | + # Rewriting unchanged state also renews its retention after a no-op request. |
| 204 | + args.output.write_text(json.dumps(state, sort_keys=True, separators=(",", ":"))) |
| 205 | + |
| 206 | + |
| 207 | +if __name__ == "__main__": |
| 208 | + main() |
0 commit comments