diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py index f02697dd80..6b347c4f83 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py @@ -14,28 +14,28 @@ class GitHubRepoSeedSource(SeedSource): seed_type: Literal["github_repo"] = "github_repo" repos: list[str] = Field( - default_factory=list, - description="List of GitHub repositories to scrape, each in `owner/name` form.", + default_factory = list, + description = "List of GitHub repositories to scrape, each in `owner/name` form.", ) token: str = Field( - default="", - description="Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.", + default = "", + description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.", ) item_types: list[Literal["issues", "pulls", "commits"]] = Field( - default=["issues", "pulls"], - description="Which GitHub item types to fetch per repo.", + default = ["issues", "pulls"], + description = "Which GitHub item types to fetch per repo.", ) limit: int = Field( - default=100, - ge=1, - le=5000, - description="Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).", + default = 100, + ge = 1, + le = 5000, + description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).", ) include_comments: bool = Field( - default=True, - description="Fetch the first N comments of each issue/PR and include them in the `comments` column.", + default = True, + description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.", ) - max_comments_per_item: int = Field(default=30, ge=0, le=200) + max_comments_per_item: int = Field(default = 30, ge = 0, le = 200) @field_validator("repos") @classmethod @@ -57,7 +57,7 @@ class GitHubRepoSeedSource(SeedSource): raise ValueError("item_types must not be empty") return list(dict.fromkeys(v)) - @model_validator(mode="after") + @model_validator(mode = "after") def _ensure_repos(self) -> "GitHubRepoSeedSource": if not self.repos: raise ValueError("At least one repo is required") diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py index e874ace527..dcd86eee7a 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py @@ -20,12 +20,12 @@ class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]): def get_dataset_uri(self) -> str: out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed" cfg = ScrapeConfig( - repos=list(self.source.repos), - token=self.source.token, - item_types=list(self.source.item_types), - limit=self.source.limit, - include_comments=self.source.include_comments, - max_comments_per_item=self.source.max_comments_per_item, + repos = list(self.source.repos), + token = self.source.token, + item_types = list(self.source.item_types), + limit = self.source.limit, + include_comments = self.source.include_comments, + max_comments_per_item = self.source.max_comments_per_item, ) path = materialize_to_jsonl(cfg, out_dir) return str(path) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py index 5a5f1dc94f..f87dbd0507 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py @@ -4,7 +4,7 @@ from data_designer.plugins.plugin import Plugin, PluginType github_repo_seed_plugin = Plugin( - impl_qualified_name="data_designer_github_repo_seed.impl.GitHubRepoSeedReader", - config_qualified_name="data_designer_github_repo_seed.config.GitHubRepoSeedSource", - plugin_type=PluginType.SEED_READER, + impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader", + config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource", + plugin_type = PluginType.SEED_READER, ) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py index fb81587c5a..5ec32739e0 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -9,6 +9,7 @@ After scraping, we read the per-resource JSONL shards and flatten them into a single unified JSONL with stable columns (`item_type`, `repo`, `number`, `title`, `body`, ...). """ + from __future__ import annotations import json @@ -59,7 +60,7 @@ def _resolve_token(token: str) -> str: def _read_jsonl(path: Path, max_rows: int | None = None): if not path.exists(): return - with path.open(encoding="utf-8") as f: + with path.open(encoding = "utf-8") as f: for i, line in enumerate(f): if not line.strip(): continue @@ -72,7 +73,11 @@ def _read_jsonl(path: Path, max_rows: int | None = None): def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: - labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")] + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] comments_nodes = (r.get("comments") or {}).get("nodes") or [] comments_text = "" if include_comments and comments_nodes: @@ -98,7 +103,11 @@ def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) - def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: - labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")] + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] comments_nodes = (r.get("comments") or {}).get("nodes") or [] comments_text = "" if include_comments and comments_nodes: @@ -126,7 +135,7 @@ def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> d def _flatten_commit_row(r: dict, repo: str) -> dict: msg = r.get("messageHeadline") or r.get("message") or "" body = r.get("messageBody") or r.get("message") or msg - author = (r.get("author") or {}) + author = r.get("author") or {} return { "item_type": "commit", "repo": repo, @@ -148,7 +157,7 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): os.environ["GH_TOKEN"] = token GitHubClient, RepoScraper = _load_impl() client = GitHubClient() - base_dir.mkdir(parents=True, exist_ok=True) + base_dir.mkdir(parents = True, exist_ok = True) # Per-resource trial limits. limit <= 0 means "all": use a very large cap. effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000 @@ -164,11 +173,11 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): for repo in cfg.repos: owner, name = repo.split("/", 1) scraper = RepoScraper( - owner=owner, - name=name, - base_dir=base_dir, - client=client, - trial_limits=trial_limits, + owner = owner, + name = name, + base_dir = base_dir, + client = client, + trial_limits = trial_limits, ) try: scraper.scrape_repo_meta() @@ -185,10 +194,18 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): repo_dir = base_dir / f"{owner}__{name}" if "issues" in cfg.item_types: for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap): - all_rows.append(_flatten_issue_row(row, repo, cfg.include_comments, cfg.max_comments_per_item)) + all_rows.append( + _flatten_issue_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) if "pulls" in cfg.item_types: for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap): - all_rows.append(_flatten_pr_row(row, repo, cfg.include_comments, cfg.max_comments_per_item)) + all_rows.append( + _flatten_pr_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) if "commits" in cfg.item_types: for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap): all_rows.append(_flatten_commit_row(row, repo)) @@ -197,13 +214,13 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path: - out_dir.mkdir(parents=True, exist_ok=True) + out_dir.mkdir(parents = True, exist_ok = True) tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120] kinds = "-".join(cfg.item_types) fname = f"github_{tag}__{kinds}__{cfg.limit}_{int(time.time())}.jsonl" out = out_dir / fname rows = scrape(cfg, out_dir / "raw") - with out.open("w", encoding="utf-8") as f: + with out.open("w", encoding = "utf-8") as f: for r in rows: - f.write(json.dumps(r, ensure_ascii=False) + "\n") + f.write(json.dumps(r, ensure_ascii = False) + "\n") return out diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py index cde03891db..bd975b3f34 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -2,6 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support.""" + from __future__ import annotations import json @@ -71,8 +72,12 @@ class GitHubClient: else: self.rest_remaining = None - def graphql(self, query: str, variables: Optional[Dict[str, Any]] = None, - max_retries: int = 20) -> Dict[str, Any]: + def graphql( + self, + query: str, + variables: Optional[Dict[str, Any]] = None, + max_retries: int = 20, + ) -> Dict[str, Any]: self._check_rate_and_wait("graphql") backoff = 2 last_err = None @@ -80,8 +85,8 @@ class GitHubClient: try: r = self.session.post( GRAPHQL_URL, - json={"query": query, "variables": variables or {}}, - timeout=120, + json = {"query": query, "variables": variables or {}}, + timeout = 120, ) self.calls_graphql += 1 # Update rate info from response headers @@ -123,7 +128,9 @@ class GitHubClient: # Retry on RATE_LIMITED for e in errs: if e.get("type") == "RATE_LIMITED": - self._sleep_until((self.graphql_reset or int(time.time()) + 60)) + self._sleep_until( + (self.graphql_reset or int(time.time()) + 60) + ) break else: # No rate-limit error, log and return partial @@ -138,8 +145,14 @@ class GitHubClient: backoff = min(backoff * 2, 60) raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}") - def rest(self, method: str, path: str, params: Optional[Dict[str, Any]] = None, - json_body: Optional[Dict[str, Any]] = None, max_retries: int = 6) -> requests.Response: + def rest( + self, + method: str, + path: str, + params: Optional[Dict[str, Any]] = None, + json_body: Optional[Dict[str, Any]] = None, + max_retries: int = 6, + ) -> requests.Response: self._check_rate_and_wait("rest") if path.startswith("http"): url = path @@ -149,7 +162,9 @@ class GitHubClient: last_err = None for attempt in range(max_retries): try: - r = self.session.request(method, url, params=params, json=json_body, timeout=120) + r = self.session.request( + method, url, params = params, json = json_body, timeout = 120 + ) self.calls_rest += 1 rem = r.headers.get("X-RateLimit-Remaining") rst = r.headers.get("X-RateLimit-Reset") @@ -190,15 +205,18 @@ class GitHubClient: backoff = min(backoff * 2, 60) raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}") - def rest_paginate(self, path: str, params: Optional[Dict[str, Any]] = None, - per_page: int = 100) -> Iterator[dict]: + def rest_paginate( + self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100 + ) -> Iterator[dict]: params = dict(params or {}) params.setdefault("per_page", per_page) url = path while True: - r = self.rest("GET", url, params=params if url == path else None) + r = self.rest("GET", url, params = params if url == path else None) if r.status_code != 200: - log.error("REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]) + log.error( + "REST paginate got %s at %s: %s", r.status_code, url, r.text[:200] + ) return items = r.json() if isinstance(items, dict): diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py index c774a5b976..b21744a885 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py @@ -131,7 +131,9 @@ def _q(parts: list[str], body: str) -> str: return "\n".join(parts + [body]) -ISSUES_PAGE_QUERY = _q([F_ACTOR, F_LABEL, F_TIMELINE], """ +ISSUES_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_TIMELINE], + """ query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) { repository(owner: $owner, name: $name) { issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { @@ -169,9 +171,12 @@ query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) } rateLimit { cost remaining resetAt } } -""") +""", +) -PRS_PAGE_QUERY = _q([F_ACTOR, F_LABEL, F_PR_TIMELINE], """ +PRS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_PR_TIMELINE], + """ query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) { repository(owner: $owner, name: $name) { pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { @@ -265,9 +270,12 @@ query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) { } rateLimit { cost remaining resetAt } } -""") +""", +) -ISSUE_COMMENTS_QUERY = _q([F_ACTOR], """ +ISSUE_COMMENTS_QUERY = _q( + [F_ACTOR], + """ query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) { repository(owner: $owner, name: $name) { issueOrPullRequest(number: $number) { @@ -298,9 +306,12 @@ query IssueComments($owner: String!, $name: String!, $number: Int!, $after: Stri } rateLimit { cost remaining resetAt } } -""") +""", +) -ISSUE_TIMELINE_QUERY = _q([F_ACTOR, F_TIMELINE], """ +ISSUE_TIMELINE_QUERY = _q( + [F_ACTOR, F_TIMELINE], + """ query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { repository(owner: $owner, name: $name) { issue(number: $number) { @@ -312,9 +323,12 @@ query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: Stri } rateLimit { cost remaining resetAt } } -""") +""", +) -PR_TIMELINE_QUERY = _q([F_ACTOR, F_PR_TIMELINE], """ +PR_TIMELINE_QUERY = _q( + [F_ACTOR, F_PR_TIMELINE], + """ query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { repository(owner: $owner, name: $name) { pullRequest(number: $number) { @@ -326,7 +340,8 @@ query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) } rateLimit { cost remaining resetAt } } -""") +""", +) PR_COMMITS_QUERY = """ query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) { @@ -364,7 +379,9 @@ query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) { } """ -PR_REVIEW_THREADS_QUERY = _q([F_ACTOR], """ +PR_REVIEW_THREADS_QUERY = _q( + [F_ACTOR], + """ query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) { repository(owner: $owner, name: $name) { pullRequest(number: $number) { @@ -389,9 +406,12 @@ query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: St } rateLimit { cost remaining resetAt } } -""") +""", +) -DISCUSSIONS_PAGE_QUERY = _q([F_ACTOR, F_LABEL], """ +DISCUSSIONS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL], + """ query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) { repository(owner: $owner, name: $name) { discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { @@ -437,9 +457,12 @@ query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: Str } rateLimit { cost remaining resetAt } } -""") +""", +) -DISCUSSION_COMMENTS_QUERY = _q([F_ACTOR], """ +DISCUSSION_COMMENTS_QUERY = _q( + [F_ACTOR], + """ query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) { repository(owner: $owner, name: $name) { discussion(number: $number) { @@ -467,9 +490,12 @@ query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: } rateLimit { cost remaining resetAt } } -""") +""", +) -DISCUSSION_REPLIES_QUERY = _q([F_ACTOR], """ +DISCUSSION_REPLIES_QUERY = _q( + [F_ACTOR], + """ query DiscussionReplies($commentId: ID!, $after: String) { node(id: $commentId) { ... on DiscussionComment { @@ -486,7 +512,8 @@ query DiscussionReplies($commentId: ID!, $after: String) { } rateLimit { cost remaining resetAt } } -""") +""", +) COMMITS_PAGE_QUERY = """ query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) { @@ -519,7 +546,9 @@ query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, } """ -RELEASES_QUERY = _q([F_ACTOR], """ +RELEASES_QUERY = _q( + [F_ACTOR], + """ query Releases($owner: String!, $name: String!, $first: Int!, $after: String) { repository(owner: $owner, name: $name) { releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { @@ -540,9 +569,12 @@ query Releases($owner: String!, $name: String!, $first: Int!, $after: String) { } rateLimit { cost remaining resetAt } } -""") +""", +) -LABELS_QUERY = _q([F_LABEL], """ +LABELS_QUERY = _q( + [F_LABEL], + """ query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) { repository(owner: $owner, name: $name) { labels(first: $first, after: $after) { @@ -552,7 +584,8 @@ query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) } rateLimit { cost remaining resetAt } } -""") +""", +) MILESTONES_QUERY = """ query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) { diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py index e6b1be866e..3695ff1641 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -5,6 +5,7 @@ Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl. """ + from __future__ import annotations import argparse @@ -34,25 +35,42 @@ def ts() -> str: class RepoScraper: - def __init__(self, owner: str, name: str, base_dir: Path, client: GitHubClient, - trial_limits: Optional[Dict[str, int]] = None): + def __init__( + self, + owner: str, + name: str, + base_dir: Path, + client: GitHubClient, + trial_limits: Optional[Dict[str, int]] = None, + ): self.owner = owner self.name = name self.base_dir = base_dir self.client = client self.trial_limits = trial_limits or {} self.repo_dir = base_dir / f"{owner}__{name}" - self.repo_dir.mkdir(parents=True, exist_ok=True) + self.repo_dir.mkdir(parents = True, exist_ok = True) self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json") # Writers self.writers: Dict[str, JsonlWriter] = {} - for key in ("issues", "pull_requests", "discussions", - "commits", "releases", "labels", "milestones", - "pr_extra_comments", "pr_extra_timeline", "pr_extra_reviews", - "issue_extra_comments", "issue_extra_timeline", - "discussion_extra_comments", "discussion_extra_replies", - "repo_meta"): + for key in ( + "issues", + "pull_requests", + "discussions", + "commits", + "releases", + "labels", + "milestones", + "pr_extra_comments", + "pr_extra_timeline", + "pr_extra_reviews", + "issue_extra_comments", + "issue_extra_timeline", + "discussion_extra_comments", + "discussion_extra_replies", + "repo_meta", + ): self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl") # ----- helpers ----- @@ -63,14 +81,25 @@ class RepoScraper: return counter >= lim def _log_rate(self, where: str, data: Dict[str, Any]) -> None: - rl = data.get("data", {}).get("rateLimit") if isinstance(data.get("data"), dict) else None + rl = ( + data.get("data", {}).get("rateLimit") + if isinstance(data.get("data"), dict) + else None + ) if rl: - log.debug("[%s] rate cost=%s remaining=%s resetAt=%s", - where, rl.get("cost"), rl.get("remaining"), rl.get("resetAt")) + log.debug( + "[%s] rate cost=%s remaining=%s resetAt=%s", + where, + rl.get("cost"), + rl.get("remaining"), + rl.get("resetAt"), + ) # ----- repo meta ----- def scrape_repo_meta(self) -> None: - data = self.client.graphql(Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}) + data = self.client.graphql( + Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name} + ) self._log_rate("repo_meta", data) repo = data.get("data", {}).get("repository") or {} repo["_fetchedAt"] = ts() @@ -89,11 +118,16 @@ class RepoScraper: per_page = 15 # conservative for heavy nested query while True: page += 1 - vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor} + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } data = self.client.graphql(Q.ISSUES_PAGE_QUERY, vars_) self._log_rate("issues", data) repo = (data.get("data") or {}).get("repository") or {} - issues = (repo.get("issues") or {}) + issues = repo.get("issues") or {} nodes = issues.get("nodes") or [] for it in nodes: it["_owner"] = self.owner @@ -101,17 +135,27 @@ class RepoScraper: it["_fetchedAt"] = ts() # Paginate nested comments/timeline if more exist if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_issue_comments(it["number"], it["comments"]["pageInfo"]["endCursor"]) + self._paginate_issue_comments( + it["number"], it["comments"]["pageInfo"]["endCursor"] + ) if it.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_issue_timeline(it["number"], it["timelineItems"]["pageInfo"]["endCursor"]) + self._paginate_issue_timeline( + it["number"], it["timelineItems"]["pageInfo"]["endCursor"] + ) if self.writers[key].write(it): total_new += 1 info = issues.get("pageInfo") or {} cursor = info.get("endCursor") self.state.set(f"{key}_cursor", cursor) - log.info("[%s/%s] issues page %d (+%d) cursor=%s remaining=%s", - self.owner, self.name, page, len(nodes), str(cursor)[:20], - self.client.graphql_remaining) + log.info( + "[%s/%s] issues page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) if self._trial_stop(key, total_new): log.info("Trial limit reached for issues (%d)", total_new) return total_new @@ -123,10 +167,17 @@ class RepoScraper: def _paginate_issue_comments(self, number: int, after: str) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("issueOrPullRequest") or {}) - comments = (item.get("comments") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} for c in comments.get("nodes") or []: c["_owner"] = self.owner c["_repo"] = self.name @@ -138,10 +189,15 @@ class RepoScraper: def _paginate_issue_timeline(self, number: int, after: str) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("issue") or {}) - tl = (item.get("timelineItems") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {} + tl = item.get("timelineItems") or {} for ev in tl.get("nodes") or []: ev["_owner"] = self.owner ev["_repo"] = self.name @@ -163,11 +219,16 @@ class RepoScraper: per_page = 3 # PR query is heavy; keep small so huge PRs don't OOM GraphQL while True: page += 1 - vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor} + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } data = self.client.graphql(Q.PRS_PAGE_QUERY, vars_) self._log_rate("prs", data) repo = (data.get("data") or {}).get("repository") or {} - prs = (repo.get("pullRequests") or {}) + prs = repo.get("pullRequests") or {} nodes = prs.get("nodes") or [] for pr in nodes: pr["_owner"] = self.owner @@ -175,23 +236,37 @@ class RepoScraper: pr["_fetchedAt"] = ts() num = pr["number"] if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_comments(num, pr["comments"]["pageInfo"]["endCursor"]) + self._paginate_pr_comments( + num, pr["comments"]["pageInfo"]["endCursor"] + ) if pr.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_timeline(num, pr["timelineItems"]["pageInfo"]["endCursor"]) + self._paginate_pr_timeline( + num, pr["timelineItems"]["pageInfo"]["endCursor"] + ) if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_commits(num, pr["commits"]["pageInfo"]["endCursor"]) + self._paginate_pr_commits( + num, pr["commits"]["pageInfo"]["endCursor"] + ) if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"): self._paginate_pr_files(num, pr["files"]["pageInfo"]["endCursor"]) if pr.get("reviewThreads", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_review_threads(num, pr["reviewThreads"]["pageInfo"]["endCursor"]) + self._paginate_pr_review_threads( + num, pr["reviewThreads"]["pageInfo"]["endCursor"] + ) if self.writers[key].write(pr): total_new += 1 info = prs.get("pageInfo") or {} cursor = info.get("endCursor") self.state.set(f"{key}_cursor", cursor) - log.info("[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s", - self.owner, self.name, page, len(nodes), str(cursor)[:20], - self.client.graphql_remaining) + log.info( + "[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) if self._trial_stop(key, total_new): log.info("Trial limit reached for PRs (%d)", total_new) return total_new @@ -203,10 +278,17 @@ class RepoScraper: def _paginate_pr_comments(self, number: int, after: str) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("issueOrPullRequest") or {}) - comments = (item.get("comments") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} for c in comments.get("nodes") or []: c["_owner"] = self.owner c["_repo"] = self.name @@ -218,10 +300,17 @@ class RepoScraper: def _paginate_pr_timeline(self, number: int, after: str) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}) - tl = (item.get("timelineItems") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + tl = item.get("timelineItems") or {} for ev in tl.get("nodes") or []: ev["_owner"] = self.owner ev["_repo"] = self.name @@ -236,10 +325,17 @@ class RepoScraper: if out_key not in self.writers: self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}) - cc = (item.get("commits") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + cc = item.get("commits") or {} for c in cc.get("nodes") or []: c["_owner"] = self.owner c["_repo"] = self.name @@ -254,10 +350,17 @@ class RepoScraper: if out_key not in self.writers: self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.PR_FILES_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}) - ff = (item.get("files") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + ff = item.get("files") or {} for f in ff.get("nodes") or []: f["_owner"] = self.owner f["_repo"] = self.name @@ -274,10 +377,17 @@ class RepoScraper: if out_key not in self.writers: self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_) - item = (((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}) - rt = (item.get("reviewThreads") or {}) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + rt = item.get("reviewThreads") or {} for th in rt.get("nodes") or []: th["_owner"] = self.owner th["_repo"] = self.name @@ -299,11 +409,16 @@ class RepoScraper: per_page = 15 while True: page += 1 - vars_ = {"owner": self.owner, "name": self.name, "first": per_page, "after": cursor} + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_) self._log_rate("discussions", data) repo = (data.get("data") or {}).get("repository") or {} - dd = (repo.get("discussions") or {}) + dd = repo.get("discussions") or {} nodes = dd.get("nodes") or [] for d in nodes: d["_owner"] = self.owner @@ -311,19 +426,29 @@ class RepoScraper: d["_fetchedAt"] = ts() num = d["number"] if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_discussion_comments(num, d["comments"]["pageInfo"]["endCursor"]) + self._paginate_discussion_comments( + num, d["comments"]["pageInfo"]["endCursor"] + ) # paginate replies per comment if needed for c in d.get("comments", {}).get("nodes", []) or []: if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_discussion_replies(c["id"], c["replies"]["pageInfo"]["endCursor"], num) + self._paginate_discussion_replies( + c["id"], c["replies"]["pageInfo"]["endCursor"], num + ) if self.writers[key].write(d): total_new += 1 info = dd.get("pageInfo") or {} cursor = info.get("endCursor") self.state.set(f"{key}_cursor", cursor) - log.info("[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s", - self.owner, self.name, page, len(nodes), str(cursor)[:20], - self.client.graphql_remaining) + log.info( + "[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) if self._trial_stop(key, total_new): return total_new if not info.get("hasNextPage"): @@ -334,10 +459,17 @@ class RepoScraper: def _paginate_discussion_comments(self, number: int, after: str) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "number": number, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_) - disc = (((data.get("data") or {}).get("repository") or {}).get("discussion") or {}) - cc = (disc.get("comments") or {}) + disc = ((data.get("data") or {}).get("repository") or {}).get( + "discussion" + ) or {} + cc = disc.get("comments") or {} for c in cc.get("nodes") or []: c["_owner"] = self.owner c["_repo"] = self.name @@ -346,10 +478,17 @@ class RepoScraper: info = cc.get("pageInfo") or {} cur = info.get("endCursor") if info.get("hasNextPage") else None - def _paginate_discussion_replies(self, comment_id: str, after: str, disc_number: int) -> None: + def _paginate_discussion_replies( + self, comment_id: str, after: str, disc_number: int + ) -> None: cur = after while cur: - vars_ = {"owner": self.owner, "name": self.name, "commentId": comment_id, "after": cur} + vars_ = { + "owner": self.owner, + "name": self.name, + "commentId": comment_id, + "after": cur, + } data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_) node = (data.get("data") or {}).get("node") or {} replies = node.get("replies") or {} @@ -374,11 +513,16 @@ class RepoScraper: per_page = 100 while True: page += 1 - vars_ = {"owner": self.owner, "name": self.name, - "first": per_page, "after": cursor, "branch": branch} + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + "branch": branch, + } data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_) self._log_rate("commits", data) - ref = (((data.get("data") or {}).get("repository") or {}).get("ref") or {}) + ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {} tgt = ref.get("target") or {} hist = tgt.get("history") or {} nodes = hist.get("nodes") or [] @@ -391,8 +535,14 @@ class RepoScraper: info = hist.get("pageInfo") or {} cursor = info.get("endCursor") self.state.set(f"{key}_cursor", cursor) - log.info("[%s/%s] commits page %d (+%d) remaining=%s", - self.owner, self.name, page, len(nodes), self.client.graphql_remaining) + log.info( + "[%s/%s] commits page %d (+%d) remaining=%s", + self.owner, + self.name, + page, + len(nodes), + self.client.graphql_remaining, + ) if self._trial_stop(key, total_new): return total_new if not info.get("hasNextPage"): @@ -417,10 +567,15 @@ class RepoScraper: return 0 total_new = 0 while True: - vars_ = {"owner": self.owner, "name": self.name, "first": 50, "after": cursor} + vars_ = { + "owner": self.owner, + "name": self.name, + "first": 50, + "after": cursor, + } data = self.client.graphql(query, vars_) repo = (data.get("data") or {}).get("repository") or {} - col = (repo.get(field) or {}) + col = repo.get(field) or {} for it in col.get("nodes") or []: it["_owner"] = self.owner it["_repo"] = self.name @@ -447,47 +602,70 @@ class RepoScraper: def setup_logging(log_file: Path) -> None: - log_file.parent.mkdir(parents=True, exist_ok=True) + log_file.parent.mkdir(parents = True, exist_ok = True) fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s" handlers = [ logging.StreamHandler(sys.stdout), - logging.FileHandler(log_file, mode="a", encoding="utf-8"), + logging.FileHandler(log_file, mode = "a", encoding = "utf-8"), ] - logging.basicConfig(level=logging.INFO, format=fmt, handlers=handlers, force=True) + logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True) def main(): ap = argparse.ArgumentParser() - ap.add_argument("--base-dir", default="/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper") - ap.add_argument("--repos", nargs="+", default=["unslothai/unsloth", "unslothai/unsloth-zoo"]) - ap.add_argument("--trial", action="store_true", help="Small trial run") - ap.add_argument("--only", nargs="+", default=None, - help="Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta") - ap.add_argument("--hf-upload-interval", type=int, default=900, - help="Seconds between HF uploads (0 to disable)") + ap.add_argument( + "--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper" + ) + ap.add_argument( + "--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"] + ) + ap.add_argument("--trial", action = "store_true", help = "Small trial run") + ap.add_argument( + "--only", + nargs = "+", + default = None, + help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta", + ) + ap.add_argument( + "--hf-upload-interval", + type = int, + default = 900, + help = "Seconds between HF uploads (0 to disable)", + ) args = ap.parse_args() base = Path(args.base_dir) data_dir = base / "data" - data_dir.mkdir(parents=True, exist_ok=True) + data_dir.mkdir(parents = True, exist_ok = True) setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log") log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial) - client = GitHubClient(min_remaining_graphql=80, min_remaining_rest=80) + client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80) rl = client.rate_snapshot() - log.info("Rate limit snapshot: %s", json.dumps(rl.get("resources", {}), default=str)[:400]) + log.info( + "Rate limit snapshot: %s", + json.dumps(rl.get("resources", {}), default = str)[:400], + ) # Start HF uploader in background if requested uploader = None if args.hf_upload_interval > 0: from hf_uploader import HFUploader - uploader = HFUploader(data_dir, interval_s=args.hf_upload_interval) + + uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval) uploader.start() trial_limits = None if args.trial: - trial_limits = {"issues": 5, "pull_requests": 5, "discussions": 3, - "commits": 20, "releases": 3, "labels": 20, "milestones": 20} + trial_limits = { + "issues": 5, + "pull_requests": 5, + "discussions": 3, + "commits": 20, + "releases": 3, + "labels": 20, + "milestones": 20, + } only = set(args.only or []) @@ -517,9 +695,12 @@ def main(): finally: if uploader: log.info("Stopping uploader and final sync...") - uploader.stop(final_upload=True) - log.info("Scraper complete. GraphQL calls=%d REST calls=%d", - client.calls_graphql, client.calls_rest) + uploader.stop(final_upload = True) + log.info( + "Scraper complete. GraphQL calls=%d REST calls=%d", + client.calls_graphql, + client.calls_rest, + ) if __name__ == "__main__": diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py index 10e836a205..efa663db2f 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py @@ -2,6 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Checkpoint state management for resumable scraping.""" + from __future__ import annotations import json @@ -14,7 +15,7 @@ from typing import Any, Dict class StateStore: def __init__(self, path: str | Path): self.path = Path(path) - self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.parent.mkdir(parents = True, exist_ok = True) self._lock = threading.Lock() self._data: Dict[str, Any] = {} if self.path.exists(): @@ -47,17 +48,18 @@ class StateStore: def _flush(self) -> None: tmp = self.path.with_suffix(self.path.suffix + ".tmp") with tmp.open("w") as f: - json.dump(self._data, f, indent=2, default=str) + json.dump(self._data, f, indent = 2, default = str) os.replace(tmp, self.path) class JsonlWriter: """Append-only JSONL writer, thread-safe, with line buffering.""" + def __init__(self, path: str | Path): self.path = Path(path) - self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.parent.mkdir(parents = True, exist_ok = True) self._lock = threading.Lock() - self._fh = self.path.open("a", buffering=1) + self._fh = self.path.open("a", buffering = 1) self._count_seen_keys: set[str] = set() # Preload seen keys if file exists (for dedup across resumes) if self.path.exists() and self.path.stat().st_size > 0: @@ -91,7 +93,7 @@ class JsonlWriter: return False if k is not None: self._count_seen_keys.add(k) - self._fh.write(json.dumps(obj, default=str, ensure_ascii=False)) + self._fh.write(json.dumps(obj, default = str, ensure_ascii = False)) self._fh.write("\n") self._fh.flush() return True