Files
quantconnect--lean/find_datasource_repos.py
T
2026-07-13 13:02:50 +08:00

246 lines
9.0 KiB
Python

#!/usr/bin/env python3
"""
Discover QuantConnect/Lean.DataSource.* repositories (non-archived) that define
the data types listed in the alternative-data dump JSON, and print them as a
semicolon-separated list on stdout.
Intended for CI: the output is captured and exported as $ADDITIONAL_STUBS_REPOS
so ci_build_stubs.sh can read it on its line 23.
Auth: REST API via urllib, reading QC_GIT_TOKEN (same env var ci_build_stubs.sh
uses when embedding the token into git clone URLs). Only 'repo' scope is
required, unlike the 'read:org' that gh's GraphQL path demanded.
"""
import json
import os
import re
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from http.client import HTTPException, HTTPMessage, HTTPSConnection
from pathlib import Path
from threading import local
from urllib.error import HTTPError
from urllib.parse import urlencode
from urllib.request import urlopen
JSON_URL = "https://s3.amazonaws.com/cdn.quantconnect.com/web/docs/alternative-data-dump-v2024-01-02.json"
OWNER = "QuantConnect"
REPO_PREFIX = "Lean.DataSource."
GH_HOST = "api.github.com"
_tls = local()
def log(msg: str) -> None:
print(msg, file=sys.stderr, flush=True)
def _conn() -> HTTPSConnection:
"""One persistent HTTPS connection per thread so repeat requests share
the TLS handshake (saves ~300-500ms per reused request)."""
conn = getattr(_tls, "conn", None)
if conn is None:
conn = HTTPSConnection(GH_HOST, timeout=30)
_tls.conn = conn
return conn
def _request(path: str, params: dict | None = None) -> tuple[object, dict[str, str]]:
"""GET a GitHub REST API path. Returns (parsed JSON, response headers)."""
url = path + ("?" + urlencode(params) if params else "")
token = os.environ.get("QC_GIT_TOKEN")
headers = {
"Accept": "application/vnd.github+json",
"X-GitHub-Api-Version": "2022-11-28",
"User-Agent": "find-datasource-repos",
"Connection": "keep-alive",
**({"Authorization": f"Bearer {token}"} if token else {}),
}
def _once() -> tuple[int, bytes, dict[str, str]]:
conn = _conn()
conn.request("GET", url, headers=headers)
resp = conn.getresponse()
return resp.status, resp.read(), dict(resp.getheaders())
try:
status, raw, resp_headers = _once()
except (HTTPException, ConnectionError, OSError):
# Stale connection; reset and retry once.
try:
_conn().close()
finally:
_tls.conn = None
status, raw, resp_headers = _once()
if status >= 400:
msg = HTTPMessage()
for k, v in resp_headers.items():
msg[k] = v
raise HTTPError(url, status, f"HTTP {status}", msg, None)
return json.loads(raw), resp_headers
def _paginate(path: str, params: dict) -> list:
"""Follow REST pagination via the Link header."""
items: list = []
page = 1
while True:
body, headers = _request(path, {**params, "per_page": 100, "page": page})
assert isinstance(body, list), f"{path} did not return a list"
items.extend(body)
if 'rel="next"' not in headers.get("Link", ""):
return items
page += 1
def fetch_data_types() -> set[str]:
log(f"Downloading {JSON_URL} ...")
with urlopen(JSON_URL) as r:
data = json.loads(r.read())
pat = re.compile(r'data-tree="QuantConnect\.DataSource\.([^"]+)"')
types: set[str] = set()
for item in data:
for doc in item.get("documentation", []):
if isinstance(doc, dict) and doc.get("title") == "Data Point Attributes":
types.update(pat.findall(doc.get("content", "")))
log(f"Extracted {len(types)} data types from 'Data Point Attributes'")
return types
def list_candidate_repos() -> dict[str, dict]:
"""Return {repo_name: {pushedAt, branch}} for all non-archived
Lean.DataSource.* repos via GET /orgs/{org}/repos."""
raw = _paginate(f"/orgs/{OWNER}/repos", {"type": "all"})
repos = {
r["name"]: {"pushedAt": r["pushed_at"], "branch": r["default_branch"]}
for r in raw
if r["name"].startswith(REPO_PREFIX) and not r.get("archived")
}
log(f"Found {len(repos)} non-archived {REPO_PREFIX}* repos")
return repos
def list_cs_files(repo: str, branch: str) -> list[str]:
"""Every .cs file path in the repo (recursive tree listing)."""
body, _ = _request(f"/repos/{OWNER}/{repo}/git/trees/{branch}", {"recursive": "1"})
assert isinstance(body, dict)
return [e["path"] for e in body.get("tree", []) if e.get("path", "").endswith(".cs")]
def search_class_definition(data_type: str) -> set[str]:
"""Fallback: REST code search for '<kw> <DataType>' in the org's .cs files."""
found: set[str] = set()
short = data_type.split(".")[-1]
for kw in ("class", "struct", "enum", "interface", "record"):
q = f'"{kw} {short}" user:{OWNER} extension:cs'
try:
body, _ = _request("/search/code", {"q": q, "per_page": 30})
except HTTPError as e:
log(f" search '{kw} {short}' failed: HTTP {e.code}")
continue
assert isinstance(body, dict)
for item in body.get("items", []):
found.add(item["repository"]["full_name"].split("/")[-1])
return found
def build_repo_to_files(repos: dict[str, dict]) -> dict[str, list[str]]:
"""Fetch the .cs file list for every repo in parallel."""
def _fetch(repo: str) -> tuple[str, list[str]]:
try:
return repo, list_cs_files(repo, repos[repo]["branch"])
except Exception as e:
log(f" warning on {repo}: {e}")
return repo, []
repo_to_files: dict[str, list[str]] = {}
total = len(repos)
with ThreadPoolExecutor(max_workers=32) as pool:
futures = [pool.submit(_fetch, r) for r in repos]
for i, fut in enumerate(as_completed(futures), 1):
repo, files = fut.result()
repo_to_files[repo] = files
log(f" [{i}/{total}] {repo} ({len(files)} .cs files)")
return repo_to_files
def match_types_to_repos(
data_types: set[str],
repo_to_files: dict[str, list[str]],
) -> tuple[dict[str, set[str]], set[str]]:
"""Return (type -> matching repos, set of unmatched types).
Primary match: a file named <DataType>.cs. Dotted types (e.g. EODHD.Events)
are nested inside outer types that get matched via a sibling non-dotted type
in the same repo, so they are skipped here."""
repo_to_basenames = {
repo: {Path(f).name for f in files} for repo, files in repo_to_files.items()
}
type_to_repos: dict[str, set[str]] = {}
unmatched: set[str] = set()
for dt in sorted(data_types):
if "." in dt:
continue
filename = f"{dt}.cs"
hits = {repo for repo, bases in repo_to_basenames.items() if filename in bases}
if hits:
type_to_repos[dt] = hits
else:
unmatched.add(dt)
return type_to_repos, unmatched
def main() -> int:
# Independent network calls run concurrently: 2.5MB JSON from S3 and
# GET /orgs/QuantConnect/repos.
with ThreadPoolExecutor(max_workers=2) as pool:
f_types = pool.submit(fetch_data_types)
f_repos = pool.submit(list_candidate_repos)
data_types = f_types.result()
repos = f_repos.result()
repo_to_files = build_repo_to_files(repos)
type_to_repos, unmatched = match_types_to_repos(data_types, repo_to_files)
if unmatched:
log(f"\n{len(unmatched)} types without filename match; "
"falling back to code search:")
valid = set(repos)
for dt in sorted(unmatched):
log(f" searching '{dt}' ...")
hits = search_class_definition(dt) & valid
if hits:
type_to_repos[dt] = hits
else:
log(f" (still no match for {dt})")
# When a type is defined in more than one repo, keep the one with the
# latest push. Example: Estimize* types live in both Lean.DataSource.Estimize
# and Lean.DataSource.ExtractAlpha; the latter is the more recent fork.
pushed_at = lambda r: repos.get(r, {}).get("pushedAt", "") # noqa: E731
picked = {dt: max(hits, key=pushed_at) for dt, hits in type_to_repos.items()}
log("\n=== Summary ===")
for dt in sorted(type_to_repos):
hits = type_to_repos[dt]
winner = picked[dt]
if len(hits) > 1:
losers = sorted(hits - {winner})
log(f" {dt:45s} -> {winner} (over: {', '.join(losers)})")
else:
log(f" {dt:45s} -> {winner}")
# ci_build_stubs.sh line 36 does ${REPO//github.com/"${TOKEN}@github.com"}
# so the clone URL must contain "github.com" -- emit full URLs, not bare names.
matching_urls = {f"https://github.com/{OWNER}/{name}" for name in picked.values()}
log(f"\n{len(matching_urls)} unique repos with matching data types")
print(";".join(sorted(matching_urls)))
if not os.environ.get("QC_GIT_TOKEN"):
log("QC_GIT_TOKEN env is empty")
return 0
if __name__ == "__main__":
sys.exit(main())