scripts/pull_gh.py

"""Pull issues + PRs + comments from a GitHub repo into JSONL.

Usage:
    python scripts/pull_gh.py --repo huggingface/peft --months 6 --out data/raw

Env: set GITHUB_TOKEN for 5000 req/hr (else 60 req/hr unauth).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

import requests

API = "https://api.github.com"


def gh_get(url: str, token: str | None, params: dict | None = None) -> dict | list:
    headers = {"Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    for attempt in range(5):
        r = requests.get(url, headers=headers, params=params, timeout=30)
        if r.status_code == 200:
            return r.json()
        if r.status_code in (403, 429):
            reset = int(r.headers.get("X-RateLimit-Reset", "0"))
            wait = max(5, reset - int(time.time()) + 2) if reset else 30
            print(f"  rate-limited; sleeping {wait}s", flush=True)
            time.sleep(min(wait, 90))
            continue
        if r.status_code == 404:
            return []
        print(f"  HTTP {r.status_code}: {r.text[:200]}", flush=True)
        time.sleep(2 + attempt)
    raise RuntimeError(f"failed after 5 attempts: {url}")


def paginate(url: str, token: str | None, params: dict, page_cap: int = 50) -> list:
    out: list = []
    page = 1
    while page <= page_cap:
        p = dict(params, page=page, per_page=100)
        chunk = gh_get(url, token, p)
        if not chunk:
            break
        out.extend(chunk)
        if len(chunk) < 100:
            break
        page += 1
    return out


def pull_repo(repo: str, since_iso: str, token: str | None, out_dir: Path) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    print(f"[pull] {repo} since {since_iso}", flush=True)

    issues = paginate(
        f"{API}/repos/{repo}/issues",
        token,
        {"state": "all", "since": since_iso, "sort": "created", "direction": "desc"},
    )
    print(f"  fetched {len(issues)} issues+prs", flush=True)
    (out_dir / "issues.jsonl").write_text(
        "\n".join(json.dumps(x) for x in issues), encoding="utf-8"
    )

    labels = gh_get(f"{API}/repos/{repo}/labels", token, {"per_page": 100})
    if isinstance(labels, list):
        (out_dir / "labels.jsonl").write_text(
            "\n".join(json.dumps(x) for x in labels), encoding="utf-8"
        )
        print(f"  fetched {len(labels)} labels", flush=True)

    comments_count = 0
    with open(out_dir / "comments.jsonl", "w", encoding="utf-8") as cf:
        for i, it in enumerate(issues):
            if it.get("comments", 0) == 0:
                continue
            num = it["number"]
            cs = paginate(
                f"{API}/repos/{repo}/issues/{num}/comments",
                token,
                {},
                page_cap=5,
            )
            for c in cs:
                c["_issue_number"] = num
                cf.write(json.dumps(c) + "\n")
                comments_count += 1
            if i % 25 == 0:
                print(f"  {i}/{len(issues)} issues, {comments_count} comments so far", flush=True)
    print(f"  fetched {comments_count} comments total", flush=True)

    return {"issues": len(issues), "comments": comments_count}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="huggingface/peft")
    ap.add_argument("--months", type=int, default=6)
    ap.add_argument("--out", default="data/raw")
    args = ap.parse_args()

    token = os.environ.get("GITHUB_TOKEN")
    if not token:
        print("WARN: no GITHUB_TOKEN; rate limit = 60/hr. Set token for 5000/hr.", flush=True)

    since = (datetime.now(timezone.utc) - timedelta(days=30 * args.months)).strftime("%Y-%m-%dT%H:%M:%SZ")
    out_dir = Path(args.out) / args.repo.replace("/", "_")
    stats = pull_repo(args.repo, since, token, out_dir)
    print(f"DONE: {stats}", flush=True)


if __name__ == "__main__":
    sys.exit(main())