spb/company-atlas

1"""`catlas` seed commands: `seed`, `import-companies`, `company-add`, `registry-stats` (docs/SEEDS.md)."""2from __future__ import annotations34from collections import Counter5from pathlib import Path6from typing import Annotated78import typer9from rich.table import Table101112def _print_counters(counters: dict[str, int]) -> None:13    from companyatlas.cli import out1415    table = Table(title="seed", show_header=False)16    for k, v in counters.items():17        table.add_row(k, str(v))18    out.print(table)192021def register(app: typer.Typer) -> None:22    @app.command("seed")23    def seed_cmd(24        no_companies: Annotated[bool, typer.Option("--no-companies", help="only industries + countries")] = False,25        limit: Annotated[int | None, typer.Option("--limit", help="stop after N registry rows")] = None,26        file: Annotated[list[Path] | None, typer.Option("--file", help="NDJSON file(s) instead of registry/companies/*.ndjson")] = None,27    ) -> None:28        """Load the seed registry (idempotent): industries, countries, companies + discover queue."""29        from companyatlas.cli import run_async30        from companyatlas.db import transaction31        from companyatlas.registry.seed import seed3233        async def go() -> dict[str, int]:34            async with transaction() as conn:35                return await seed(conn, companies=not no_companies, limit=limit, files=file or None)3637        _print_counters(run_async(go()))3839    @app.command("import-companies")40    def import_cmd(path: Annotated[Path, typer.Argument(help=".ndjson / .jsonl / .json / .csv with a `website` column")],41                   source: Annotated[str, typer.Option("--source")] = "manual") -> None:42        """Import companies from a manual file (`website` required; `display_name` derived from the domain when missing)."""43        from companyatlas.cli import run_async44        from companyatlas.db import transaction45        from companyatlas.registry.seed import import_companies, read_rows_file4647        rows = read_rows_file(path)4849        async def go() -> dict[str, int]:50            async with transaction() as conn:51                return await import_companies(conn, rows, source=source)5253        _print_counters(run_async(go()))5455    @app.command("company-add")56    def add_cmd(website: Annotated[str, typer.Argument()],57                name: Annotated[str | None, typer.Option("--name")] = None,58                country: Annotated[str | None, typer.Option("--country", help="ISO-2")] = None,59                industry: Annotated[list[str] | None, typer.Option("--industry", help="taxonomy slug (repeatable)")] = None) -> None:60        """Add one company by website and queue its discovery."""61        from companyatlas.cli import out, run_async62        from companyatlas.db import transaction63        from companyatlas.registry.seed import add_company6465        async def go() -> dict:66            async with transaction() as conn:67                return await add_company(conn, website, display_name=name, country=country, industries=industry or [])6869        row = run_async(go())70        counters = row.pop("counters", {})71        out.print(row)72        _print_counters(counters)7374    @app.command("registry-stats")75    def stats_cmd(top: Annotated[int, typer.Option("--top", help="rows per table")] = 30) -> None:76        """Counts by country / industry / tier from the registry NDJSON files (no database)."""77        from companyatlas.cli import out78        from companyatlas.registry.industries import industry_index, top_level_of79        from companyatlas.registry.seed import load_countries_file, load_registry_rows, registry_files8081        rows = load_registry_rows()82        countries = {c["code"]: c for c in load_countries_file()}83        names = {slug: ind.name for slug, ind in industry_index().items()}84        out.print(f"[bold]{len(rows)}[/] companies in {len(registry_files())} file(s); "85                  f"{sum(1 for r in rows if r.get('public_company'))} public; {sum(1 for r in rows if not r.get('industries'))} without industry")8687        def table(title: str, counter: Counter, label=lambda k: k) -> None:  # type: ignore[no-untyped-def]88            t = Table(title=title)89            t.add_column("key")90            t.add_column("companies", justify="right")91            t.add_column("share", justify="right")92            for k, v in counter.most_common(top):93                t.add_row(str(label(k)), str(v), f"{100 * v / max(1, len(rows)):.1f} %")94            out.print(t)9596        table("tiers", Counter(r.get("tier") for r in rows))97        table("regions", Counter(countries.get(r.get("country") or "", {}).get("region") or "unknown" for r in rows))98        table("countries", Counter(r.get("country") or "??" for r in rows), lambda k: f"{k} {countries.get(k, {}).get('name', '')}".strip())99        table("top-level industries", Counter(t for r in rows for t in {top_level_of(s) for s in r.get("industries", [])}), lambda k: names.get(k, k))100        table("industries", Counter(s for r in rows for s in r.get("industries", [])), lambda k: names.get(k, k))101