""" data_loader.py — loads current grant snapshots (JSON) and optional past winners This module keeps IO concerns simple and robust: - Recursively loads current-grant JSON files under a snapshots directory - Optionally loads past winners from either an Excel file or a JSON directory - Returns Python lists of dictionaries; no model code here Public API --------- load_current_grants(snapshots_dir: Path | str, limit: int | None = None) -> list[dict] load_past_winners(history_xlsx: Path | str | None = None, history_json_dir: Path | str | None = None) -> list[dict] """ from __future__ import annotations from pathlib import Path from typing import Any, Dict, List, Optional import json import logging import pandas as pd logger = logging.getLogger(__name__) # ----------------------------- Current grants --------------------------------- def load_current_grants(snapshots_dir: Path | str, limit: Optional[int] = None) -> List[Dict[str, Any]]: """Load current grant JSON snapshots from a directory tree. Each file is expected to be one JSON object. The function tolerates missing keys and will attach an `id` from the filename if not present. """ snapshots_dir = Path(snapshots_dir) if not snapshots_dir.exists(): logger.warning("snapshots directory not found: %s", snapshots_dir) return [] records: List[Dict[str, Any]] = [] for p in sorted(snapshots_dir.rglob("*.json")): try: with open(p, "r", encoding="utf-8") as f: rec = json.load(f) if not isinstance(rec, dict): logger.debug("Skipping non-object JSON: %s", p) continue rec.setdefault("id", p.stem) rec.setdefault("_path", str(p)) records.append(rec) if limit and len(records) >= limit: break except Exception as e: # pragma: no cover logger.warning("Failed to load %s: %s", p, e) continue logger.info("Loaded %d current grants from %s", len(records), snapshots_dir) return records # ------------------------------ Past winners ---------------------------------- _CANON_COLS = { # canonical : candidate column names (lower/underscore) "project_title": ["project_title", "title", "name"], "abstract": ["abstract", "description", "summary", "public_description"], "competition": ["competition", "programme", "program"], "award_amount": ["award_amount", "amount", "grant", "project_cost", "award"], "lead_org": ["lead_org", "lead_organisation", "lead_organization", "organisation_name", "organization_name"], "year": ["year", "fy", "start_year"], "project_url": ["project_url", "url", "link"], } def _norm_cols(df: pd.DataFrame) -> pd.DataFrame: df = df.copy() df.columns = [ (c if isinstance(c, str) else str(c)) .lower() .replace(" ", "_") .replace("-", "_") for c in df.columns ] return df def _ensure_canonical(df: pd.DataFrame) -> pd.DataFrame: for canon, candidates in _CANON_COLS.items(): if canon in df.columns: continue for c in candidates: if c in df.columns: df[canon] = df[c] break if canon not in df.columns: df[canon] = None return df def _load_past_winners_from_excel(xlsx_path: Path) -> List[Dict[str, Any]]: df = pd.read_excel(xlsx_path) df = _norm_cols(df) df = _ensure_canonical(df) return [row._asdict() if hasattr(row, "_asdict") else row.to_dict() for _, row in df.iterrows()] def _load_past_winners_from_json_dir(json_dir: Path) -> List[Dict[str, Any]]: records: List[Dict[str, Any]] = [] for p in sorted(json_dir.rglob("*.json")): try: with open(p, "r", encoding="utf-8") as f: rec = json.load(f) if not isinstance(rec, dict): continue rec.setdefault("_path", str(p)) records.append(rec) except Exception: continue return records def _load_past_winners_from_jsonl(jsonl_path: Path) -> List[Dict[str, Any]]: """Load past winners from JSONL file (optionally gzipped).""" import gzip records: List[Dict[str, Any]] = [] # Check if file is gzipped open_func = gzip.open if str(jsonl_path).endswith('.gz') else open try: with open_func(jsonl_path, 'rt', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue try: rec = json.loads(line) if isinstance(rec, dict): records.append(rec) except json.JSONDecodeError: continue return records except Exception as e: logger.warning(f"Failed to load JSONL from {jsonl_path}: {e}") return [] def load_past_winners( history_xlsx: Path | str | None = None, history_json_dir: Path | str | None = None, history_jsonl: Path | str | None = None ) -> List[Dict[str, Any]]: """Load past winners from Excel, JSONL, or JSON directory. Priority order: JSONL > JSON dir > Excel. If neither exists, returns an empty list. """ # Try JSONL first (most efficient for large datasets) if history_jsonl is not None: jsonl = Path(history_jsonl) if jsonl.exists(): recs = _load_past_winners_from_jsonl(jsonl) if recs: logger.info("Loaded %d past winners from JSONL: %s", len(recs), jsonl) return recs # Also check for past_winners.jsonl.gz in default location (for HF deployment) default_jsonl = Path("data/past_winners.jsonl.gz") if default_jsonl.exists() and history_jsonl is None: recs = _load_past_winners_from_jsonl(default_jsonl) if recs: logger.info("Loaded %d past winners from default JSONL: %s", len(recs), default_jsonl) return recs # JSON dir next if history_json_dir is not None: jdir = Path(history_json_dir) if jdir.exists(): recs = _load_past_winners_from_json_dir(jdir) if recs: logger.info("Loaded %d past winners from JSON dir: %s", len(recs), jdir) return recs else: logger.info("No JSON past winners found under %s", jdir) # Excel last if history_xlsx is not None: xlsx = Path(history_xlsx) if xlsx.exists(): recs = _load_past_winners_from_excel(xlsx) logger.info("Loaded %d past winners from Excel: %s", len(recs), xlsx) return recs else: logger.info("History Excel not found: %s", xlsx) return [] # Self-test if __name__ == "__main__": import argparse logging.basicConfig(level=logging.INFO) parser = argparse.ArgumentParser() parser.add_argument("--snapshots-dir", type=Path, default=Path("data/snapshots")) parser.add_argument("--history-xlsx", type=Path, default=Path("data/past_winners.xlsx")) parser.add_argument("--limit", type=int, default=3) args = parser.parse_args() current = load_current_grants(args.snapshots_dir, limit=args.limit) history = load_past_winners(args.history_xlsx) print(f"current: {len(current)} | history: {len(history)}") if current: print("example current keys:", sorted(current[0].keys())[:12]) if history: print("example history keys:", sorted(history[0].keys())[:12])