evalstate/diffusers-pr-api

1from __future__ import annotations2 3import json4import time5import urllib.error6import urllib.parse7import urllib.request8from collections.abc import Callable, Iterator9from datetime import UTC, datetime10 11from slop_farmer.data.http import urlopen_with_retry12 13 14class PullRequestDiffTooLargeError(RuntimeError):15    """Raised when GitHub refuses to render a PR diff because it is too large."""16 17 18class GitHubApiRequestError(RuntimeError):19    """Raised when GitHub returns a non-recoverable HTTP response."""20 21    def __init__(self, status_code: int, path: str, detail: str):22        self.status_code = status_code23        self.path = path24        self.detail = detail25        super().__init__(f"GitHub API request failed: {status_code} {path} {detail}")26 27 28class GitHubClient:29    def __init__(30        self,31        token: str | None,32        per_page: int = 100,33        timeout: int = 180,34        max_retries: int = 5,35        log: Callable[[str], None] | None = None,36    ):37        self.token = token38        self.per_page = per_page39        self.timeout = timeout40        self.max_retries = max_retries41        self.log = log42        self.base_url = "https://api.github.com"43        self.request_count = 044        self.rate_limit_log_every = 2545        self._rate_limit_thresholds_logged: dict[str, set[int]] = {}46 47    @staticmethod48    def _header_int(headers: dict[str, str], name: str) -> int | None:49        value = headers.get(name)50        if value is None:51            return None52        try:53            return int(value)54        except ValueError:55            return None56 57    def _maybe_log_rate_limit(self, path: str, headers: dict[str, str]) -> None:58        self.request_count += 159        if not self.log:60            return61        limit = self._header_int(headers, "x-ratelimit-limit")62        remaining = self._header_int(headers, "x-ratelimit-remaining")63        used = self._header_int(headers, "x-ratelimit-used")64        reset_at = self._header_int(headers, "x-ratelimit-reset")65        resource = headers.get("x-ratelimit-resource", "?")66        if limit is None or remaining is None:67            return68        should_log = self.request_count == 1 or self.request_count % self.rate_limit_log_every == 069        thresholds = (1000, 500, 250, 100, 50, 25, 10, 5, 1, 0)70        logged = self._rate_limit_thresholds_logged.setdefault(resource, set())71        for threshold in thresholds:72            if remaining <= threshold and threshold not in logged:73                should_log = True74                logged.add(threshold)75        if not should_log:76            return77        reset_text = "?"78        if reset_at is not None:79            reset_text = datetime.fromtimestamp(reset_at, tz=UTC).strftime("%Y-%m-%dT%H:%M:%SZ")80        used_text = "?" if used is None else str(used)81        self.log(82            f"GitHub rate limit: resource={resource} used={used_text} remaining={remaining}/{limit} "83            f"reset={reset_text} after {path} (request {self.request_count})"84        )85 86    def _request(87        self,88        path: str,89        params: dict[str, object] | None = None,90        accept: str = "application/vnd.github+json",91    ) -> tuple[object, dict[str, str]]:92        query = f"?{urllib.parse.urlencode(params)}" if params else ""93        url = f"{self.base_url}{path}{query}"94        request = urllib.request.Request(url)95        request.add_header("Accept", accept)96        request.add_header("X-GitHub-Api-Version", "2022-11-28")97        if self.token:98            request.add_header("Authorization", f"Bearer {self.token}")99        while True:100            try:101                with urlopen_with_retry(102                    request,103                    timeout=self.timeout,104                    max_retries=self.max_retries,105                    log=self.log,106                    label=path,107                ) as response:108                    payload = response.read().decode("utf-8")109                    headers = {k.lower(): v for k, v in response.headers.items()}110                    self._maybe_log_rate_limit(path, headers)111                    return json.loads(payload), headers112            except urllib.error.HTTPError as exc:113                if exc.code == 403 and exc.headers.get("X-RateLimit-Remaining") == "0":114                    reset_at = int(exc.headers.get("X-RateLimit-Reset", "0") or "0")115                    sleep_for = max(reset_at - int(time.time()), 1)116                    if self.log:117                        self.log(118                            f"GitHub rate limit reached for {path}; sleeping {sleep_for}s until reset"119                        )120                    time.sleep(sleep_for)121                    continue122                detail = exc.read().decode("utf-8", errors="replace")123                if (124                    exc.code == 406125                    and '"field":"diff"' in detail126                    and '"code":"too_large"' in detail127                ):128                    raise PullRequestDiffTooLargeError(129                        f"GitHub diff too large: {path} {detail}"130                    ) from exc131                raise GitHubApiRequestError(exc.code, path, detail) from exc132 133    def _request_text(134        self,135        path: str,136        params: dict[str, object] | None = None,137        accept: str = "application/vnd.github.diff",138    ) -> tuple[str, dict[str, str]]:139        query = f"?{urllib.parse.urlencode(params)}" if params else ""140        url = f"{self.base_url}{path}{query}"141        request = urllib.request.Request(url)142        request.add_header("Accept", accept)143        request.add_header("X-GitHub-Api-Version", "2022-11-28")144        if self.token:145            request.add_header("Authorization", f"Bearer {self.token}")146        while True:147            try:148                with urlopen_with_retry(149                    request,150                    timeout=self.timeout,151                    max_retries=self.max_retries,152                    log=self.log,153                    label=path,154                ) as response:155                    payload = response.read().decode("utf-8", errors="replace")156                    headers = {k.lower(): v for k, v in response.headers.items()}157                    self._maybe_log_rate_limit(path, headers)158                    return payload, headers159            except urllib.error.HTTPError as exc:160                if exc.code == 403 and exc.headers.get("X-RateLimit-Remaining") == "0":161                    reset_at = int(exc.headers.get("X-RateLimit-Reset", "0") or "0")162                    sleep_for = max(reset_at - int(time.time()), 1)163                    if self.log:164                        self.log(165                            f"GitHub rate limit reached for {path}; sleeping {sleep_for}s until reset"166                        )167                    time.sleep(sleep_for)168                    continue169                detail = exc.read().decode("utf-8", errors="replace")170                if (171                    exc.code == 406172                    and '"field":"diff"' in detail173                    and '"code":"too_large"' in detail174                ):175                    raise PullRequestDiffTooLargeError(176                        f"GitHub diff too large: {path} {detail}"177                    ) from exc178                raise GitHubApiRequestError(exc.code, path, detail) from exc179 180    def paginate(181        self,182        path: str,183        params: dict[str, object] | None = None,184        accept: str = "application/vnd.github+json",185        limit: int | None = None,186    ) -> Iterator[dict]:187        page = 1188        yielded = 0189        params = dict(params or {})190        params["per_page"] = self.per_page191        while True:192            params["page"] = page193            payload, _headers = self._request(path, params=params, accept=accept)194            if not isinstance(payload, list):195                raise RuntimeError(f"Expected list payload from {path}, got {type(payload)!r}")196            if not payload:197                break198            for item in payload:199                if not isinstance(item, dict):200                    continue201                yield item202                yielded += 1203                if limit is not None and yielded >= limit:204                    return205            if len(payload) < self.per_page:206                break207            page += 1208 209    def get_json(self, path: str, accept: str = "application/vnd.github+json") -> dict:210        payload, _headers = self._request(path, accept=accept)211        if not isinstance(payload, dict):212            raise RuntimeError(f"Expected dict payload from {path}, got {type(payload)!r}")213        return payload214 215    def iter_repo_issues(216        self, owner: str, repo: str, since: str | None, limit: int | None217    ) -> Iterator[dict]:218        direction = "asc"219        if since is None and limit is not None:220            # When the caller asks for a bounded first pass without a watermark,221            # prefer the most recently updated items. Using ascending order here222            # causes small smoke tests to read the stalest issue/PR stubs first,223            # which can make recent-age filters appear to return an empty repo.224            direction = "desc"225        params: dict[str, object] = {"state": "all", "sort": "updated", "direction": direction}226        if since:227            params["since"] = since228        yield from self.paginate(f"/repos/{owner}/{repo}/issues", params=params, limit=limit)229 230    def iter_issue_comments(231        self, owner: str, repo: str, since: str | None, limit: int | None232    ) -> Iterator[dict]:233        params: dict[str, object] = {"sort": "updated", "direction": "asc"}234        if since:235            params["since"] = since236        yield from self.paginate(237            f"/repos/{owner}/{repo}/issues/comments", params=params, limit=limit238        )239 240    def iter_issue_comments_for_number(241        self,242        owner: str,243        repo: str,244        number: int,245        since: str | None,246        limit: int | None = None,247    ) -> Iterator[dict]:248        params: dict[str, object] = {"sort": "updated", "direction": "asc"}249        if since:250            params["since"] = since251        yield from self.paginate(252            f"/repos/{owner}/{repo}/issues/{number}/comments", params=params, limit=limit253        )254 255    def get_pull_request(self, owner: str, repo: str, number: int) -> dict:256        return self.get_json(f"/repos/{owner}/{repo}/pulls/{number}")257 258    def iter_pull_reviews(259        self, owner: str, repo: str, number: int, limit: int | None = None260    ) -> Iterator[dict]:261        yield from self.paginate(f"/repos/{owner}/{repo}/pulls/{number}/reviews", limit=limit)262 263    def iter_pull_review_comments(264        self, owner: str, repo: str, number: int, limit: int | None = None265    ) -> Iterator[dict]:266        yield from self.paginate(267            f"/repos/{owner}/{repo}/pulls/{number}/comments",268            params={"sort": "updated", "direction": "asc"},269            limit=limit,270        )271 272    def iter_pull_files(273        self, owner: str, repo: str, number: int, limit: int | None = None274    ) -> Iterator[dict]:275        yield from self.paginate(f"/repos/{owner}/{repo}/pulls/{number}/files", limit=limit)276 277    def get_pull_request_diff(self, owner: str, repo: str, number: int) -> str:278        path = f"/repos/{owner}/{repo}/pulls/{number}"279        try:280            payload, _headers = self._request_text(281                path,282                accept="application/vnd.github.diff",283            )284            return payload285        except PullRequestDiffTooLargeError:286            if self.log:287                self.log(288                    f"Skipping unified diff for pull request #{number}; GitHub reports diff too large"289                )290            return ""291 292    def iter_issue_timeline(293        self, owner: str, repo: str, number: int, limit: int | None = None294    ) -> Iterator[dict]:295        path = f"/repos/{owner}/{repo}/issues/{number}/timeline"296        try:297            yield from self.paginate(298                path,299                accept="application/vnd.github+json, application/vnd.github.mockingbird-preview+json",300                limit=limit,301            )302        except GitHubApiRequestError as exc:303            if exc.status_code < 500:304                raise305            if self.log:306                self.log(307                    f"Skipping timeline fetch for issue #{number} after GitHub {exc.status_code}: {path}"308                )309            return310