File size: 7,612 Bytes
59ebe66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf9b3dc
 
 
59ebe66
cf9b3dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59ebe66
 
cf9b3dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59ebe66
 
 
 
 
 
 
 
 
 
cf9b3dc
59ebe66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
"""
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])