stack/dev/scripts/generate_cms_express.py

"""Generate CMS express modules from table model introspection.

Reads every ``src/cms/table/*.py`` model and emits one
``@nw.narwhalify`` function per table grouped into domain modules
under ``src/cms/express/``.

Each generated function calls ``auto_cast(df)`` which inspects
column names at runtime and casts string columns matching CMS
numeric naming conventions to Float64.

Usage::

    uv run python dev/scripts/generate_cms_express.py
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

# ── Constants ────────────────────────────────────────────────────

SRC = Path("src")
TABLE_DIR = SRC / "cms" / "table"
EXPRESS_DIR = SRC / "cms" / "express"

# ── Domain grouping ──────────────────────────────────────────────

DOMAIN_RULES: list[tuple[str, list[str]]] = [
    (
        "aco",
        [
            "accountable_care_organization",
            "aco_reach_",
            "number_accountable_care",
            "county_level_aggregate_expenditure",
            "performance_year_financial",
            "reach_acos",
            "advance_investment_payment",
            "pioneer_aco",
            "value_modifier",
        ],
    ),
    (
        "enrollment",
        [
            "medicare_monthly_enrollment",
            "program_statistics_medicare_total_enrollment",
            "program_statistics_original_medicare_enrollment",
            "program_statistics_medicare_advantage_other_health_plan_enrollment",
            "program_statistics_medicare_part_d_enrollment",
            "program_statistics_medicare_newly_enrolled",
            "program_statistics_medicare_medicaid_dual_enrollment",
            "program_statistics_medicare_deaths",
            "program_statistics_medicare_premiums",
            "medicare_fee_service_public_provider_enrollment",
        ],
    ),
    (
        "provider",
        [
            "medicare_physician_other_practitioners",
            "medicare_part_d_prescribers",
            "physician_supplier_procedure",
            "order_referring",
            "opt_out_affidavits",
            "fiscal_intermediary_shared_system",
            "pending_initial_logging_tracking",
            "revalidation_",
            "provider_services_file_",
            "public_reporting_missing",
            "managing_clinician_",
            "quality_payment_program",
            "restructured_betos",
            "medicare_provider_supplier_taxonomy",
            "medicare_clinical_laboratory",
            "medicare_fee_service_comprehensive",
        ],
    ),
    (
        "facility",
        [
            "hospital_",
            "skilled_nursing_facility_",
            "home_health_agency_",
            "hospice_",
            "federally_qualified_health_center_",
            "rural_health_clinic_",
            "long_term_care_facility",
            "facility_level_minimum_data_set",
            "minimum_data_set_frequency",
            "nursing_home_chain",
            "payroll_based_journal",
            "deficit_reduction_act",
            "end_stage_renal_disease",
            "medicare_dialysis_facilities",
            "home_infusion_therapy",
            "income_asset_ownership",
            "opioid_treatment_program",
        ],
    ),
    (
        "drug_spending",
        [
            "medicare_part_b_spending",
            "medicare_part_b_discarded",
            "medicare_part_d_spending",
            "medicare_quarterly_part_b",
            "medicare_quarterly_part_d",
            "medicaid_spending_by_drug",
            "medicaid_opioid_",
            "medicare_part_d_opioid",
            "monthly_prescription_drug",
            "quarterly_prescription_drug",
        ],
    ),
    (
        "utilization",
        [
            "medicare_inpatient_hospitals",
            "medicare_outpatient_hospitals",
            "medicare_durable_medical_equipment",
            "medicare_post_acute_care",
            "medicare_geographic_variation",
            "medicare_advantage_geographic_variation",
            "medicare_telehealth",
            "medicare_covid_19",
            "medicare_current_beneficiary_survey",
            "specialty_",
            "beneficiary_",
            "post_acute_care",
            "ambulatory_surgery_center",
        ],
    ),
    (
        "market",
        [
            "market_saturation",
            "medicare_advantage_",
            "medicare_plan_finder",
            "prescription_drug_plan_",
            "medicaid_managed_care",
            "medicare_demonstrations",
            "medicare_diabetes",
        ],
    ),
    (
        "innovation",
        [
            "innovation_center_",
            "cpc_initiative_",
            "comprehensive_care_joint_replacement",
            "kidney_care_choices",
            "ambulatory_specialty_model",
            "strong_start_",
            "agency_healthcare_research",
        ],
    ),
    (
        "program_statistics",
        [
            "program_statistics_medicare_inpatient",
            "program_statistics_medicare_outpatient",
            "program_statistics_medicare_home_health",
            "program_statistics_medicare_hospice",
            "program_statistics_medicare_skilled_nursing",
            "program_statistics_medicare_physician",
            "program_statistics_medicare_part_d",
            "program_statistics_medicare_part_part_b",
            "program_statistics_medicare_providers",
            "program_statistics_medicare_advantage_inpatient",
            "program_statistics_medicare_advantage_outpatient",
            "program_statistics_medicare_advantage_physician",
            "program_statistics_medicare_advantage_skilled",
        ],
    ),
]


def classify_domain(table_stem: str) -> str:
    """Assign a table to a domain module."""
    for domain, prefixes in DOMAIN_RULES:
        for prefix in prefixes:
            if table_stem.startswith(prefix) or prefix in table_stem:
                return domain
    return "other"


# ── Table introspection ──────────────────────────────────────────


def parse_table_file(path: Path) -> dict | None:
    """Extract class name, tablename, and field count from a table module."""
    content = path.read_text()
    class_match = re.search(r"class (\w+)\(SQLTable\):", content)
    if not class_match:
        return None
    class_name = class_match.group(1)

    tn_match = re.search(r'__tablename__\s*=\s*"([^"]+)"', content)
    tablename = tn_match.group(1) if tn_match else path.stem

    field_count = len(re.findall(r"    \w+: (?:str|date) \| None", content))

    return {
        "class_name": class_name,
        "tablename": tablename,
        "stem": path.stem,
        "field_count": field_count,
    }


# ── Code generation ──────────────────────────────────────────────


def generate_function(table: dict) -> str:
    """Generate a single @nw.narwhalify function for a table."""
    tablename = table["tablename"]
    param = f"cms__{tablename}"
    field_count = table["field_count"]

    lines = [
        "@nw.narwhalify",
        f"def {tablename}({param}: FrameT) -> FrameT:",
        f'    """Clean and type-cast cms.{tablename}.',
        "",
        f"    {field_count} fields.  Numeric columns auto-cast to Float64.",
        '    """',
        f"    return auto_cast({param})",
    ]
    return "\n".join(lines)


def generate_module(domain: str, tables: list[dict]) -> str:
    """Generate a full express module for a domain group."""
    lines = [
        f'"""CMS express — {domain} domain.',
        "",
        f"Auto-generated narwhals functions for {len(tables)} CMS tables.",
        "Each function takes the raw table and returns a typed frame",
        "with numeric columns cast to Float64.",
        '"""',
        "",
        "from __future__ import annotations",
        "",
        "import narwhals as nw",
        "from narwhals.typing import FrameT",
        "",
        "from cms.express._helpers import auto_cast",
    ]

    for tbl in sorted(tables, key=lambda t: t["tablename"]):
        lines.append("")
        lines.append("")
        lines.append(generate_function(tbl))

    lines.append("")
    return "\n".join(lines)


def generate_init(domains: dict[str, list[dict]]) -> str:
    """Generate src/cms/express/__init__.py."""
    lines = [
        '"""CMS express layer — narwhals transformation functions.',
        "",
        "Auto-generated pure functions that clean, type-cast, and",
        "standardize all 150 CMS public datasets.",
        "",
        "Each domain module provides one function per table:",
        "",
    ]
    for domain in sorted(domains):
        count = len(domains[domain])
        lines.append(f"- **{domain}**: {count} tables")
    lines.append('"""')
    lines.append("")

    for domain in sorted(domains):
        lines.append(f"from . import {domain} as {domain}")

    lines.append("")
    return "\n".join(lines)


# ── Main ─────────────────────────────────────────────────────────


def main() -> None:
    tables = []
    for f in sorted(TABLE_DIR.glob("*.py")):
        if f.name == "__init__.py":
            continue
        info = parse_table_file(f)
        if info and info["field_count"] >= 2:
            tables.append(info)

    print(f"Parsed {len(tables)} table models")

    domains: dict[str, list[dict]] = {}
    for tbl in tables:
        domain = classify_domain(tbl["stem"])
        domains.setdefault(domain, []).append(tbl)

    for domain, tbls in sorted(domains.items()):
        print(f"  {domain}: {len(tbls)} tables")

    EXPRESS_DIR.mkdir(parents=True, exist_ok=True)

    for domain, tbls in domains.items():
        module_path = EXPRESS_DIR / f"{domain}.py"
        content = generate_module(domain, tbls)
        module_path.write_text(content)
        print(f"  wrote {module_path} ({len(tbls)} functions)")

    init_path = EXPRESS_DIR / "__init__.py"
    content = generate_init(domains)
    init_path.write_text(content)
    print(f"  wrote {init_path}")

    total_fns = sum(len(v) for v in domains.values())
    print(f"\nDone: {total_fns} functions in {len(domains)} modules")


if __name__ == "__main__":
    sys.path.insert(0, str(SRC))
    main()