school_compare/pipeline/scripts/generate_gias_codes.py

"""Generate GIAS code->name dictionaries from the live bulk CSV.

Writes:
  - backend/gias_codes.py                        (canonical Python module)
  - pipeline/scripts/gias_codes.py               (byte-identical copy)
  - pipeline/transform/seeds/gias_code_names.csv (dbt seed for drift test)

Run from the repo root whenever the dbt drift test warns that DfE
added/renamed a value:  python pipeline/scripts/generate_gias_codes.py
"""

from __future__ import annotations

import io
import sys
from datetime import date, timedelta
from pathlib import Path

import pandas as pd
import requests

GIAS_URL = (
    "https://ea-edubase-api-prod.azurewebsites.net"
    "/edubase/downloads/public/edubasealldata{date}.csv"
)

# (CSV code column, CSV name column, python dict name, seed field key)
FIELDS = [
    ("TypeOfEstablishment (code)", "TypeOfEstablishment (name)", "SCHOOL_TYPE", "school_type"),
    ("EstablishmentStatus (code)", "EstablishmentStatus (name)", "ESTABLISHMENT_STATUS", "establishment_status"),
    ("PhaseOfEducation (code)", "PhaseOfEducation (name)", "PHASE_OF_EDUCATION", "phase_of_education"),
    ("OfficialSixthForm (code)", "OfficialSixthForm (name)", "OFFICIAL_SIXTH_FORM", "official_sixth_form"),
    ("ReligiousCharacter (code)", "ReligiousCharacter (name)", "RELIGIOUS_CHARACTER", "religious_character"),
    ("AdmissionsPolicy (code)", "AdmissionsPolicy (name)", "ADMISSIONS_POLICY", "admissions_policy"),
]

MODULE_HEADER = '''"""GIAS code -> name dictionaries.

GENERATED by pipeline/scripts/generate_gias_codes.py from the GIAS bulk CSV
— do not edit by hand; rerun the script when the dbt drift test warns.
The canonical file is backend/gias_codes.py; pipeline/scripts/gias_codes.py
must be byte-identical (enforced by backend/tests/test_gias_codes.py).
"""

from __future__ import annotations

import logging
import math

logger = logging.getLogger(__name__)

'''

MODULE_FOOTER = '''

def translate(code, mapping: dict[int, str]) -> str | None:
    """Translate a GIAS code to its display name.

    None/NaN -> None (column absent or suppressed). Unknown codes degrade to
    "Unknown (<code>)" with a warning so a new DfE value never blanks the UI.
    """
    if code is None or (isinstance(code, float) and math.isnan(code)):
        return None
    code = int(code)
    if code not in mapping:
        logger.warning("Unknown GIAS code %s (not in dictionary)", code)
        return f"Unknown ({code})"
    return mapping[code]
'''


def download_csv() -> pd.DataFrame:
    for day in (date.today(), date.today() - timedelta(days=1)):
        url = GIAS_URL.format(date=day.strftime("%Y%m%d"))
        print(f"Downloading {url}")
        resp = requests.get(url, timeout=300)
        if resp.status_code == 404:
            continue
        resp.raise_for_status()
        return pd.read_csv(
            io.StringIO(resp.content.decode("latin-1")),
            dtype=str, keep_default_na=False,
        )
    sys.exit("GIAS CSV not available for today or yesterday")


def main() -> None:
    repo = Path(__file__).resolve().parents[2]
    df = download_csv()

    module_parts = [MODULE_HEADER]
    seed_rows: list[tuple[str, int, str]] = []

    for code_col, name_col, dict_name, field_key in FIELDS:
        pairs = (
            df[[code_col, name_col]]
            .loc[lambda d: d[code_col] != ""]
            .drop_duplicates()
        )
        by_code: dict[int, set] = {}
        for c, n in pairs.itertuples(index=False):
            by_code.setdefault(int(c), set()).add(n)
        mapping = []
        for code, names in sorted(by_code.items()):
            named = sorted(n for n in names if n != "")
            if len(named) > 1:
                sys.exit(f"{code_col}: code {code} maps to multiple names {named} — investigate before generating")
            # Codes that only ever appear with a blank (name) are GIAS
            # "not recorded" sentinels (e.g. ReligiousCharacter 99,
            # AdmissionsPolicy 9). Map them to "" so the API serves the same
            # empty string the old name pipeline did — the "Unknown (<code>)"
            # path is reserved for genuinely new codes.
            mapping.append((code, named[0] if named else ""))
        lines = [f"{dict_name}: dict[int, str] = {{"]
        for code, name in mapping:
            escaped = name.replace('"', '\\"')
            lines.append(f'    {code}: "{escaped}",')
        lines.append("}\n")
        module_parts.append("\n".join(lines))
        seed_rows += [(field_key, code, name) for code, name in mapping]

    module = "\n".join(module_parts) + MODULE_FOOTER

    (repo / "backend" / "gias_codes.py").write_text(module)
    (repo / "pipeline" / "scripts" / "gias_codes.py").write_text(module)

    seed_path = repo / "pipeline" / "transform" / "seeds" / "gias_code_names.csv"
    with open(seed_path, "w", newline="") as fh:
        import csv
        w = csv.writer(fh)
        w.writerow(["field", "code", "name"])
        w.writerows(seed_rows)

    print(f"Wrote backend/gias_codes.py, pipeline/scripts/gias_codes.py, {seed_path.name}")
    print("\nKey codes for the dbt work (Task 3):")
    for field in ("establishment_status", "phase_of_education", "official_sixth_form"):
        print(f"  {field}:")
        for f, code, name in seed_rows:
            if f == field:
                print(f"    {code} = {name}")


if __name__ == "__main__":
    main()