thermograph/gen_cities.py

"""Offline generator for backend/cities.json — the finite set of cities that get
crawlable climate pages (/climate/<slug>). Run occasionally to refresh the list:

    python gen_cities.py [N]        # default N=500 top metros by population

It reuses the GeoNames index that places.py already downloads/parses (calling
places._load() synchronously fills places._data), takes the top-N places by
population, and assigns each a stable, unique, URL-safe slug. Committing the output
keeps the routable city set explicit and reviewable, and decouples page-serving
from the async place-name loader.
"""
import json
import os
import re
import sys
import unicodedata

import places

OUT_PATH = os.path.join(os.path.dirname(__file__), "cities.json")

# GeoNames entry tuple layout (see places._load): the fields we keep.
_NAME, _ADMIN1, _COUNTRY, _CC, _LAT, _LON, _POP = 1, 2, 3, 4, 5, 6, 7


def slugify(*parts: str) -> str:
    """ASCII, lowercase, hyphenated slug from name/admin/country parts."""
    text = " ".join(p for p in parts if p)
    text = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode()
    text = re.sub(r"[^a-zA-Z0-9]+", "-", text).strip("-").lower()
    return re.sub(r"-{2,}", "-", text)


# Core English-speaking countries — the first English top-up tier.
ENGLISH_CC = {"US", "GB", "CA", "AU", "NZ", "IE", "ZA"}

# A second English top-up tier: the remaining countries where English is an official
# language, plus countries where >35% of the population speaks English (Eurobarometer
# 2012 "can hold a conversation in English" / EF EPI). Both drive English-language
# search demand. Editable — this is a judgment call, not a hard rule.
ENGLISH_EXTENDED_CC = {
    # English official (beyond the core seven)
    "IN", "PK", "PH", "SG", "HK", "MY", "LK", "PG", "FJ",
    "NG", "KE", "GH", "UG", "TZ", "ZW", "ZM", "MW", "BW", "NA", "RW", "SL", "LR", "MU", "SS", "SZ", "LS", "GM", "SC",
    "JM", "TT", "BB", "BS", "BZ", "GY", "GD", "LC", "VC", "AG", "DM", "KN", "MT",
    # >35% English proficiency (non-official)
    "NL", "SE", "DK", "NO", "IS", "FI", "DE", "AT", "BE", "CH", "LU", "CY", "SI", "GR", "EE", "LV", "LT", "FR", "IL",
}


def _to_city(e, seen_slugs: set[str]) -> dict | None:
    name, admin1, country, cc = e[_NAME], e[_ADMIN1], e[_COUNTRY], e[_CC]
    # Drop admin1 from the slug when it just repeats the city name
    # (e.g. Tokyo/Tokyo, Singapore/Singapore) to avoid "tokyo-tokyo-jp".
    admin_part = admin1 if admin1 and slugify(admin1) != slugify(name) else ""
    base = slugify(name, admin_part, cc or "")
    if not base:
        return None
    slug = base
    i = 2
    while slug in seen_slugs:      # disambiguate the rare collision
        slug = f"{base}-{i}"
        i += 1
    seen_slugs.add(slug)
    return {
        "slug": slug, "name": name, "admin1": admin1,
        "country": country, "country_code": cc,
        "lat": round(e[_LAT], 5), "lon": round(e[_LON], 5),
        "population": e[_POP],
    }


def build(n_global: int = 500, n_english: int = 250, n_extended: int = 250) -> list[dict]:
    """Three tiers, population-descending, de-duplicated:
      1. top n_global cities worldwide,
      2. up to n_english more from core English-speaking countries,
      3. up to n_extended more from the remaining English-official + >35%-English
         countries — all not already chosen."""
    places._load()  # synchronous parse; fills places._data (entries are pop-desc)
    if not places._data:
        raise SystemExit("GeoNames index failed to load (see logs); cannot generate cities.")
    entries = places._data[0]

    out: list[dict] = []
    seen_slugs: set[str] = set()
    chosen_ids: set = set()   # (name, admin1, cc) already added, so tiers don't overlap

    def add_from(pred, limit: int) -> int:
        added = 0
        for e in entries:
            if added >= limit:
                break
            if not pred(e):
                continue
            ident = (e[_NAME], e[_ADMIN1], e[_CC])
            if ident in chosen_ids:
                continue
            c = _to_city(e, seen_slugs)
            if c:
                out.append(c)
                chosen_ids.add(ident)
                added += 1
        return added

    add_from(lambda e: True, n_global)                              # tier 1: global
    add_from(lambda e: e[_CC] in ENGLISH_CC, n_english)             # tier 2: core English
    add_from(lambda e: e[_CC] in ENGLISH_EXTENDED_CC, n_extended)   # tier 3: extended English
    return out


def main() -> None:
    a = sys.argv[1:]
    n_global = int(a[0]) if len(a) > 0 else 500
    n_english = int(a[1]) if len(a) > 1 else 250
    n_extended = int(a[2]) if len(a) > 2 else 250
    cities = build(n_global, n_english, n_extended)
    with open(OUT_PATH, "w", encoding="utf-8") as f:
        json.dump(cities, f, ensure_ascii=False, indent=0, separators=(",", ":"))
        f.write("\n")
    print(f"wrote {len(cities)} cities ({n_global} global + {n_english} core-English + "
          f"{n_extended} extended-English) -> {OUT_PATH}")
    print("sample:", ", ".join(c["slug"] for c in cities[:8]))


if __name__ == "__main__":
    main()