scripts/build_census_cbp_dataset.py

import csv
import json
import os
import urllib.request
from pathlib import Path

OUT_DIR = Path("data/census_cbp")
OUT_DIR.mkdir(parents=True, exist_ok=True)

YEAR = "2023"

api_key = os.environ.get("CENSUS_API_KEY", "").strip()

if not api_key:
    raise SystemExit(
        "Missing Census API key. Run: export CENSUS_API_KEY='YOUR_KEY_HERE'"
    )

URL = (
    f"https://api.census.gov/data/{YEAR}/cbp"
    "?get=NAME,NAICS2017_LABEL,ESTAB,EMP,PAYANN,PAYQTR1"
    "&for=state:*"
    "&NAICS2017=00"
    "&LFO=001"
    "&EMPSZES=001"
    f"&key={api_key}"
)

print("Downloading Census County Business Patterns state-level all-sector data...")
safe_url = URL.replace(api_key, "****") if api_key else URL
print(safe_url)

with urllib.request.urlopen(URL, timeout=60) as response:
    raw = response.read().decode("utf-8", errors="replace")

try:
    rows = json.loads(raw)
except json.JSONDecodeError:
    print("Census API returned non-JSON response:")
    print(raw[:1000])
    raise SystemExit(1)

headers = rows[0]
records = []

for row in rows[1:]:
    item = dict(zip(headers, row))

    def to_int(value):
        if value in (None, "", "null"):
            return None
        return int(value)

    records.append({
        "year": int(YEAR),
        "geo_level": "state",
        "state_name": item.get("NAME"),
        "state_code": item.get("state"),
        "naics_code": item.get("NAICS2017"),
        "naics_label": item.get("NAICS2017_LABEL"),
        "establishments": to_int(item.get("ESTAB")),
        "employment": to_int(item.get("EMP")),
        "annual_payroll_1000_usd": to_int(item.get("PAYANN")),
        "first_quarter_payroll_1000_usd": to_int(item.get("PAYQTR1")),
        "source": "U.S. Census Bureau County Business Patterns API",
        "license": "U.S. Government public data"
    })

jsonl_path = OUT_DIR / "us_cbp_state_all_sectors_2023.jsonl"
csv_path = OUT_DIR / "us_cbp_state_all_sectors_2023.csv"

with jsonl_path.open("w", encoding="utf-8") as f:
    for record in records:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

with csv_path.open("w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=list(records[0].keys()))
    writer.writeheader()
    writer.writerows(records)

print(f"Done. Records: {len(records)}")
print(f"Saved: {jsonl_path}")
print(f"Saved: {csv_path}")