spb/company-atlas

1#!/usr/bin/env python2"""SEC EDGAR enrichment for the seed registry (docs/SEEDS.md).34Matches registry companies to `company_tickers.json` (ticker first, then normalised name) and fills *missing* `sec_cik`, `ticker`,5`exchange`, `hq_city`/`hq_region` and an industry hint (SIC → taxonomy slug) from the submissions JSON. Existing values are never6overwritten. Polite: SEC-required User-Agent, ≤ 5 requests/s, responses cached under data/seed/edgar/.78Usage: .venv/bin/python scripts/seed_edgar.py [--max-submissions N] [--dry-run]9"""10from __future__ import annotations1112import argparse13import json14import logging15import sys16import time17from collections import Counter, defaultdict18from pathlib import Path19from typing import Any2021import httpx2223ROOT = Path(__file__).resolve().parents[1]24sys.path.insert(0, str(ROOT / "src"))2526from companyatlas.ids import normalize_alias27from companyatlas.registry.industries import map_sic2829log = logging.getLogger("seed_edgar")30USER_AGENT = "CompanyAtlasBot/0.1 contact@spboucher.ai"31TICKERS_URL = "https://www.sec.gov/files/company_tickers.json"32SUBMISSIONS_URL = "https://data.sec.gov/submissions/CIK{cik:010d}.json"33MIN_INTERVAL_S = 0.21               # ≤ 5 requests per second (SEC fair-access policy is 10/s)34CACHE_DIR = ROOT / "data" / "seed" / "edgar"35OUT_DIR = ROOT / "registry" / "companies"36US_STATES = {"AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "FL", "GA", "HI", "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MD", "MA", "MI",37             "MN", "MS", "MO", "MT", "NE", "NV", "NH", "NJ", "NM", "NY", "NC", "ND", "OH", "OK", "OR", "PA", "RI", "SC", "SD", "TN", "TX", "UT",38             "VT", "VA", "WA", "WV", "WI", "WY", "DC", "PR"}394041class Edgar:42    def __init__(self) -> None:43        self.client = httpx.Client(headers={"User-Agent": USER_AGENT, "Accept-Encoding": "gzip, deflate"}, timeout=30, follow_redirects=True)44        self.last = 0.045        self.requests = 046        CACHE_DIR.mkdir(parents=True, exist_ok=True)4748    def get_json(self, url: str, cache_name: str) -> Any | None:49        path = CACHE_DIR / cache_name50        if path.exists():51            return json.loads(path.read_text(encoding="utf-8"))52        for attempt in range(4):53            wait = MIN_INTERVAL_S - (time.monotonic() - self.last)54            if wait > 0:55                time.sleep(wait)56            try:57                r = self.client.get(url)58            except httpx.HTTPError as e:59                log.warning("edgar %s: %s", url, e)60                time.sleep(2 * (attempt + 1))61                continue62            finally:63                self.last = time.monotonic()64                self.requests += 165            if r.status_code == 200:66                path.write_text(r.text, encoding="utf-8")67                return r.json()68            if r.status_code == 404:69                path.write_text("null", encoding="utf-8")70                return None71            log.warning("edgar %s → HTTP %d, backing off", url, r.status_code)72            time.sleep(5 * (attempt + 1))73        return None747576def load_registry() -> dict[Path, list[dict[str, Any]]]:77    out: dict[Path, list[dict[str, Any]]] = {}78    for path in sorted(OUT_DIR.glob("wikidata-*.ndjson")):79        with path.open(encoding="utf-8") as f:80            out[path] = [json.loads(line) for line in f if line.strip()]81    return out828384def main() -> int:85    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)86    ap.add_argument("--max-submissions", type=int, default=8000, help="cap on submissions JSON fetches (≈ 30 min at 5 req/s when cold)")87    ap.add_argument("--dry-run", action="store_true")88    ap.add_argument("-v", "--verbose", action="store_true")89    args = ap.parse_args()90    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(asctime)s %(levelname)s %(message)s")91    logging.getLogger("httpx").setLevel(logging.WARNING)9293    edgar = Edgar()94    tickers = edgar.get_json(TICKERS_URL, "company_tickers.json")95    if not tickers:96        log.error("could not fetch %s", TICKERS_URL)97        return 198    by_ticker: dict[str, list[dict[str, Any]]] = defaultdict(list)99    by_name: dict[str, set[int]] = defaultdict(set)100    for e in tickers.values():101        by_ticker[e["ticker"].upper()].append(e)102        by_name[normalize_alias(e["title"])].add(int(e["cik_str"]))103    log.info("edgar tickers: %d entries, %d names", len(tickers), len(by_name))104105    files = load_registry()106    stats: Counter[str] = Counter()107    submissions_fetched = 0108    for path, rows in files.items():109        for r in rows:110            cik: int | None = int(r["sec_cik"]) if r.get("sec_cik") and str(r["sec_cik"]).isdigit() else None111            if cik is None:112                cands: set[int] = set()113                if r.get("ticker") and r["ticker"].upper() in by_ticker:114                    for e in by_ticker[r["ticker"].upper()]:115                        if normalize_alias(e["title"])[:6] == normalize_alias(r["display_name"])[:6] or r.get("country") == "US":116                            cands.add(int(e["cik_str"]))117                if not cands:118                    names = [r["display_name"], r.get("legal_name") or ""] + list(r.get("aliases") or [])119                    for nm in names:120                        if nm and normalize_alias(nm) in by_name:121                            cands |= by_name[normalize_alias(nm)]122                    if len(cands) > 1 and r.get("country") != "US":123                        cands = set()      # ambiguous non-US name: do not guess124                if len(cands) == 1:125                    cik = cands.pop()126                    stats["matched_name_or_ticker"] += 1127                elif len(cands) > 1:128                    stats["ambiguous"] += 1129                    continue130                else:131                    continue132            else:133                stats["had_cik"] += 1134            r.setdefault("industry_labels", [])135            changed = False136            if not r.get("sec_cik"):137                r["sec_cik"] = str(cik)138                changed = True139            needs_sub = not r.get("exchange") or not r.get("industries") or not r.get("hq_city") or not r.get("ticker")140            if needs_sub and submissions_fetched < args.max_submissions:141                sub = edgar.get_json(SUBMISSIONS_URL.format(cik=cik), f"CIK{cik:010d}.json")142                submissions_fetched += 1143                if sub:144                    if not r.get("ticker") and sub.get("tickers"):145                        r["ticker"] = sub["tickers"][0]146                        changed = True147                    if not r.get("exchange") and sub.get("exchanges"):148                        r["exchange"] = next((x for x in sub["exchanges"] if x), None)149                        changed = changed or bool(r["exchange"])150                    if r.get("ticker"):151                        r["public_company"] = True152                    sic = sub.get("sic")153                    slug = map_sic(sic)154                    if sic and sub.get("sicDescription") and f"SIC {sic} {sub['sicDescription']}" not in r["industry_labels"]:155                        r["industry_labels"].append(f"SIC {sic} {sub['sicDescription']}")156                        changed = True157                    if slug and not r.get("industries"):158                        r["industries"] = [slug]159                        stats["industry_from_sic"] += 1160                        changed = True161                    biz = (sub.get("addresses") or {}).get("business") or {}162                    if not r.get("hq_city") and biz.get("city"):163                        r["hq_city"] = biz["city"].title()164                        changed = True165                    if not r.get("hq_region") and biz.get("stateOrCountry") in US_STATES:166                        r["hq_region"] = biz["stateOrCountry"]167                        changed = True168            if changed:169                r["source_edgar"] = {"cik": str(cik), "enriched_at": time.strftime("%Y-%m-%dT%H:%M:%S+00:00", time.gmtime())}170                stats["rows_enriched"] += 1171        if not args.dry_run:172            with path.open("w", encoding="utf-8") as f:173                for r in rows:174                    f.write(json.dumps(r, ensure_ascii=False) + "\n")175    log.info("edgar done: %s (%d HTTP requests, %d submissions)", dict(stats), edgar.requests, submissions_fetched)176    (ROOT / "data" / "seed" / "edgar-report.json").write_text(json.dumps({"stats": stats, "submissions": submissions_fetched}, indent=1))177    if not args.dry_run:178        total = sum(len(rows) for rows in files.values())179        with_cik = sum(1 for rows in files.values() for r in rows if r.get("sec_cik"))180        no_ind = sum(1 for rows in files.values() for r in rows if not r.get("industries"))181        update_readme(f"* EDGAR ({time.strftime('%Y-%m-%d', time.gmtime())}): {stats['matched_name_or_ticker']:,} companies newly matched to a CIK "182                      f"(ticker, then exact normalised name), {stats['had_cik']:,} already had one → {with_cik:,} of {total:,} with a SEC CIK; "183                      f"{stats['rows_enriched']:,} rows enriched, {stats['industry_from_sic']:,} industries from SIC ({no_ind:,} still without an "184                      f"industry mapping), {submissions_fetched:,} submissions JSON read.", stats, submissions_fetched)185    return 0186187188def update_readme(line: str, stats: Counter[str], submissions: int) -> None:189    """Replace (or append) the EDGAR line of the README harvest report and record the counters in data/seed/harvest-stats.json."""190    readme = OUT_DIR / "README.md"191    if readme.exists():192        lines = [ln for ln in readme.read_text(encoding="utf-8").splitlines() if not ln.startswith("* EDGAR (")]193        while lines and not lines[-1].strip():194            lines.pop()195        readme.write_text("\n".join([*lines, line]) + "\n", encoding="utf-8")196    stats_path = ROOT / "data" / "seed" / "harvest-stats.json"197    data = json.loads(stats_path.read_text(encoding="utf-8")) if stats_path.exists() else {}198    data["edgar"] = {**stats, "submissions": submissions, "at": time.strftime("%Y-%m-%dT%H:%M:%S+00:00", time.gmtime())}199    stats_path.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")200201202if __name__ == "__main__":203    sys.exit(main())204