Emi Griffith
b702e019d5
api_suggest carried ~80 lines of ranking policy — the exactness-boosted scoring, local/upstream blending with dedupe, and the token-respelling correction loop — all consuming the match: prefix/fuzzy vocabulary that places.py produces. Producer and consumer now live together: places.suggest(q, upstream, limit) implements the whole policy with the upstream geocoder lookup injected as a callable, so places stays free of the fetch layer and the policy is testable without FastAPI. The endpoint is the HTTP shim: call places.suggest, map the nothing-at-all-to-serve failure to a 502. Behavior unchanged. Policy tests: upstream skipped when the local answer convinces, blend + dedupe, the exactness boost (a hamlet spelled like the typo can't beat Seattle), single-typo queries answered by the fuzzy matcher without a correction, the correction path verified via the upstream probe, upstream failure degrading to local results, and the terminal error-propagation case.
344 lines
14 KiB
Python
344 lines
14 KiB
Python
"""Typo-tolerant place-name index for search suggestions.
|
||
|
||
Open-Meteo's geocoder (climate.geocode) only matches exact spellings — one
|
||
mistyped letter and the query returns nothing. This module keeps a local index
|
||
of world places (a GeoNames cities dump, downloaded once into data/geonames/
|
||
and reused across restarts) so /suggest can answer instantly and tolerate a
|
||
single-letter typo — a substituted, missing, or extra letter, or two adjacent
|
||
letters swapped — anywhere in the query, including the first character.
|
||
|
||
It also builds a vocabulary of place-name tokens so a typo'd word in a
|
||
multi-word query can be respelled against words the index knows ("pest
|
||
seattle" → "west seattle"); the endpoint verifies those candidates against the
|
||
upstream geocoder, which covers neighbourhood-level places the cities dump
|
||
lacks.
|
||
|
||
The index loads in a background thread (start_loading()). Until it's ready —
|
||
or forever, if the download fails — search()/corrections() return None/empty
|
||
and /suggest degrades to the plain upstream geocoder, so the app never needs
|
||
this data to boot or serve.
|
||
"""
|
||
import bisect
|
||
import heapq
|
||
import os
|
||
import threading
|
||
import time
|
||
import unicodedata
|
||
import zipfile
|
||
|
||
import httpx
|
||
|
||
import audit
|
||
|
||
GEO_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "geonames")
|
||
GEONAMES_URL = "https://download.geonames.org/export/dump/"
|
||
|
||
# Which GeoNames cities dump to index. cities1000 (~170k places, population
|
||
# ≥ 1000) is small enough to hold in memory and big enough to cover the tiny
|
||
# vacation towns people actually search for; set THERMOGRAPH_CITIES=cities5000
|
||
# (or cities15000) to trade coverage for a lighter footprint.
|
||
CITIES = os.environ.get("THERMOGRAPH_CITIES", "cities1000")
|
||
|
||
# Entry tuple layout: (norm_name, name, admin1, country, country_code, lat, lon, pop)
|
||
_POP = 7
|
||
|
||
_load_lock = threading.Lock()
|
||
_load_started = False
|
||
|
||
# Set once, atomically, by the loader thread: (entries, names, order, vocab)
|
||
# where `entries` is population-descending (so fuzzy scans can stop at the
|
||
# first matches found), `names`/`order` are the normalized names sorted
|
||
# alphabetically with their entry indices (for prefix bisection), and `vocab`
|
||
# maps each place-name token to the population of the biggest place using it.
|
||
_data = None
|
||
|
||
|
||
def start_loading() -> None:
|
||
"""Kick off the background index load, once per process (idempotent)."""
|
||
global _load_started
|
||
with _load_lock:
|
||
if _load_started:
|
||
return
|
||
_load_started = True
|
||
threading.Thread(target=_load, name="places-index", daemon=True).start()
|
||
|
||
|
||
def ready() -> bool:
|
||
return _data is not None
|
||
|
||
|
||
def _fetch(name: str) -> str:
|
||
"""Path to a GeoNames dump file, downloading into data/geonames/ if absent.
|
||
Cached forever — cities don't move; delete the folder to force a refresh."""
|
||
path = os.path.join(GEO_DIR, name)
|
||
if os.path.exists(path) and os.path.getsize(path) > 0:
|
||
return path
|
||
r = httpx.get(GEONAMES_URL + name, timeout=120, follow_redirects=True)
|
||
r.raise_for_status()
|
||
tmp = path + ".part"
|
||
with open(tmp, "wb") as f:
|
||
f.write(r.content)
|
||
os.replace(tmp, path) # atomic: never leave a truncated file behind
|
||
return path
|
||
|
||
|
||
def _load() -> None:
|
||
t0 = time.monotonic()
|
||
try:
|
||
os.makedirs(GEO_DIR, exist_ok=True)
|
||
# Admin-division and country display names ("US.WA" → Washington).
|
||
admin1 = {}
|
||
with open(_fetch("admin1CodesASCII.txt"), encoding="utf-8") as f:
|
||
for line in f:
|
||
cols = line.rstrip("\n").split("\t")
|
||
if len(cols) >= 2:
|
||
admin1[cols[0]] = cols[1]
|
||
countries = {}
|
||
with open(_fetch("countryInfo.txt"), encoding="utf-8") as f:
|
||
for line in f:
|
||
if line.startswith("#"):
|
||
continue
|
||
cols = line.split("\t")
|
||
if len(cols) >= 5:
|
||
countries[cols[0]] = cols[4]
|
||
entries = []
|
||
with zipfile.ZipFile(_fetch(f"{CITIES}.zip")) as z, z.open(f"{CITIES}.txt") as f:
|
||
for raw in f:
|
||
cols = raw.decode("utf-8").rstrip("\n").split("\t")
|
||
if len(cols) < 15:
|
||
continue
|
||
name, ascii_name, cc, a1 = cols[1], cols[2], cols[8], cols[10]
|
||
norm = _norm(ascii_name or name)
|
||
if len(norm) < 2:
|
||
continue
|
||
try:
|
||
lat, lon, pop = float(cols[4]), float(cols[5]), int(cols[14] or 0)
|
||
except ValueError:
|
||
continue
|
||
entries.append((norm, name, admin1.get(f"{cc}.{a1}"),
|
||
countries.get(cc), cc, lat, lon, pop))
|
||
global _data
|
||
_data = _build(entries)
|
||
except Exception as e: # noqa: BLE001 - suggestions degrade to the upstream geocoder
|
||
audit.log_event("error", {"phase": "places_load", "error": repr(e),
|
||
"seconds": round(time.monotonic() - t0, 1)})
|
||
|
||
|
||
def _build(entries: list) -> tuple:
|
||
entries.sort(key=lambda e: e[_POP], reverse=True)
|
||
order = sorted(range(len(entries)), key=lambda i: entries[i][0])
|
||
names = [entries[i][0] for i in order]
|
||
vocab = {}
|
||
for e in entries:
|
||
for tok in e[0].split():
|
||
if len(tok) >= 3 and vocab.get(tok, -1) < e[_POP]:
|
||
vocab[tok] = e[_POP]
|
||
return (entries, names, order, vocab)
|
||
|
||
|
||
def _norm(s: str) -> str:
|
||
"""Lowercased, accent-stripped, punctuation-flattened matching key, so
|
||
"Coeur d'alene" finds Cœur d'Alene and "winston salem" Winston-Salem."""
|
||
s = unicodedata.normalize("NFKD", s)
|
||
s = "".join(c for c in s if not unicodedata.combining(c))
|
||
for ch in ("'", "’", "."):
|
||
s = s.replace(ch, "")
|
||
for ch in ("-", ",", "/"):
|
||
s = s.replace(ch, " ")
|
||
return " ".join(s.casefold().split())
|
||
|
||
|
||
def _prefix_edit1(q: str, w: str) -> bool:
|
||
"""True when `q` is within one edit — a substituted, extra, or missing
|
||
letter, or an adjacent swap — of some prefix of `w`. Prefix matching (not
|
||
whole-name) so "chicgo" already suggests Chicago mid-typing."""
|
||
n = len(q)
|
||
i = 0
|
||
m = min(n, len(w))
|
||
while i < m and q[i] == w[i]:
|
||
i += 1
|
||
if i == n:
|
||
return True # exact prefix, zero edits
|
||
if w.startswith(q[i + 1:], i + 1):
|
||
return True # one letter substituted
|
||
if w.startswith(q[i + 1:], i):
|
||
return True # one extra letter typed
|
||
if w.startswith(q[i:], i + 1):
|
||
return True # one letter missed
|
||
return (i + 1 < n and i + 1 < len(w) and q[i] == w[i + 1] and q[i + 1] == w[i]
|
||
and w.startswith(q[i + 2:], i + 2)) # adjacent letters swapped
|
||
|
||
|
||
def _within1(a: str, b: str) -> bool:
|
||
"""Whole-token Damerau-Levenshtein distance ≤ 1 (used for the vocabulary)."""
|
||
la, lb = len(a), len(b)
|
||
if abs(la - lb) > 1:
|
||
return False
|
||
i = 0
|
||
m = min(la, lb)
|
||
while i < m and a[i] == b[i]:
|
||
i += 1
|
||
if la == lb:
|
||
if i == la:
|
||
return True
|
||
if a[i + 1:] == b[i + 1:]:
|
||
return True # substitution
|
||
return (i + 1 < la and a[i] == b[i + 1] and a[i + 1] == b[i]
|
||
and a[i + 2:] == b[i + 2:]) # transposition
|
||
s, l = (a, b) if la < lb else (b, a)
|
||
return s[i:] == l[i + 1:] # insertion/deletion
|
||
|
||
|
||
def _result(e: tuple, match: str) -> dict:
|
||
return {"name": e[1], "admin1": e[2], "country": e[3], "country_code": e[4],
|
||
"lat": e[5], "lon": e[6], "population": e[_POP], "match": match}
|
||
|
||
|
||
def search(q: str, limit: int = 5) -> list[dict] | None:
|
||
"""Top `limit` places matching `q` as a (possibly typo'd) name prefix.
|
||
|
||
Exact-prefix matches come first, population-descending; remaining slots are
|
||
filled with names one edit away (only for queries of 4+ characters — with
|
||
fewer, "one letter off" matches everything). Returns None while the index
|
||
isn't loaded so the caller can fall back to the upstream geocoder.
|
||
"""
|
||
data = _data
|
||
if data is None:
|
||
return None
|
||
entries, names, order, _ = data
|
||
qn = _norm(q)
|
||
if len(qn) < 2:
|
||
return []
|
||
lo = bisect.bisect_left(names, qn)
|
||
hi = bisect.bisect_left(names, qn + "\uffff", lo)
|
||
top = heapq.nlargest(limit, range(lo, hi), key=lambda i: entries[order[i]][_POP])
|
||
out = [_result(entries[order[i]], "prefix") for i in top]
|
||
if len(out) < limit and len(qn) >= 4:
|
||
need = limit - len(out)
|
||
# A single edit can only touch one of the first two characters, so a
|
||
# candidate's first two must overlap the query's — a cheap filter that
|
||
# rejects ~85% of the index before the real per-name check. Scanning in
|
||
# population order means we can stop at the first `need` hits.
|
||
q0, q1 = qn[0], qn[1]
|
||
for e in entries:
|
||
w = e[0]
|
||
if w[0] != q0 and w[0] != q1 and (len(w) < 2 or (w[1] != q0 and w[1] != q1)):
|
||
continue
|
||
if w.startswith(qn):
|
||
continue # already counted in the prefix tier
|
||
if _prefix_edit1(qn, w):
|
||
out.append(_result(e, "fuzzy"))
|
||
need -= 1
|
||
if need == 0:
|
||
break
|
||
return out
|
||
|
||
|
||
def _suggest_score(r: dict, exact: bool) -> int:
|
||
"""Prominence rank with an exactness boost — big enough that a real place
|
||
beats a same-size typo match, small enough that a hamlet spelled exactly
|
||
like the typo can't outrank a metropolis one edit away ("Seatle", a
|
||
village, must not beat Seattle)."""
|
||
pop = r.get("population") or 0
|
||
return pop * 4 + 1 if exact else pop
|
||
|
||
|
||
def _suggest_merge(local: list, upstream: tuple, limit: int) -> list:
|
||
"""Blend local-index and upstream results into one top-`limit` list.
|
||
Everything the user's spelling matches exactly counts as exact; only the
|
||
local index's typo-tolerant matches take the fuzzy (unboosted) score."""
|
||
scored = [(r, _suggest_score(r, r.get("match") != "fuzzy")) for r in local]
|
||
scored += [(r, _suggest_score(r, True)) for r in upstream]
|
||
scored.sort(key=lambda t: t[1], reverse=True)
|
||
out, seen = [], set()
|
||
for r, _ in scored:
|
||
key = ((r.get("name") or "").casefold(), r.get("admin1") or "",
|
||
r.get("country_code") or "")
|
||
if key not in seen:
|
||
seen.add(key)
|
||
out.append(r)
|
||
if len(out) == limit:
|
||
break
|
||
return out
|
||
|
||
|
||
def suggest(q: str, upstream, limit: int = 5) -> tuple[list[dict], str | None]:
|
||
"""The whole suggestion policy: the top `limit` places for a (possibly
|
||
typo'd) query prefix, plus the respelling that produced them (or None).
|
||
|
||
The local index answers instantly and tolerates a single-letter typo;
|
||
``upstream(q) -> iterable`` (an injected geocoder lookup, so this module
|
||
stays free of the fetch layer) fills remaining slots with what the index
|
||
doesn't know (neighbourhoods, postcodes). The upstream is consulted unless
|
||
the local index already answered convincingly: full slots and a substantial
|
||
place matching the spelling as typed. Fuzzy hits don't count as convincing —
|
||
they're guesses, and when the query is an alternate name the index doesn't
|
||
know ("münchen" prefix-matches only small towns; Munich lives upstream),
|
||
neither those towns nor a coincidental fuzzy big-city hit may suppress the
|
||
real answer. When both come up empty, one query token is respelled against
|
||
known place-name tokens and retried ("pest seattle" → "west seattle").
|
||
|
||
An upstream failure degrades to whatever the index found — it propagates
|
||
only when there is nothing at all to serve (index still loading, no
|
||
correction hits)."""
|
||
q = q.strip()
|
||
if len(q) < 2:
|
||
return [], None
|
||
local = search(q, limit) # None while the index loads
|
||
results = list(local or [])
|
||
upstream_error = None
|
||
prominent = max((r.get("population") or 0 for r in results
|
||
if r.get("match") == "prefix"), default=0)
|
||
if len(results) < limit or prominent < 100_000:
|
||
try:
|
||
results = _suggest_merge(results, tuple(upstream(q)), limit)
|
||
except Exception as e: # noqa: BLE001 - local results (if any) still serve
|
||
upstream_error = e
|
||
corrected = None
|
||
if not results and len(q) >= 4:
|
||
for phrase in corrections(q):
|
||
hits = search(phrase, limit) or []
|
||
if not hits:
|
||
try:
|
||
hits = list(upstream(phrase))
|
||
except Exception: # noqa: BLE001 - a failed probe just means no hits
|
||
hits = []
|
||
if hits:
|
||
results, corrected = hits[:limit], phrase
|
||
break
|
||
if not results and local is None and upstream_error is not None:
|
||
raise upstream_error
|
||
return results[:limit], corrected
|
||
|
||
|
||
def corrections(q: str, max_phrases: int = 3) -> list[str]:
|
||
"""Candidate respellings of `q` with one token replaced by a known
|
||
place-name token a single edit away — "pest seattle" → "west seattle".
|
||
|
||
Ranked by how likely the swap is: tokens the vocabulary has never seen get
|
||
corrected first, and replacements are ordered by the population of the
|
||
biggest place using them ("west" over "wesh"). Callers verify candidates by
|
||
actually searching, so a wrong guess only costs one lookup.
|
||
"""
|
||
data = _data
|
||
if data is None:
|
||
return []
|
||
vocab = data[3]
|
||
toks = _norm(q).split()
|
||
cands: dict[str, tuple] = {}
|
||
for i, t in enumerate(toks):
|
||
if len(t) < 3:
|
||
continue # respelling 1-2 letter tokens is noise
|
||
unknown = t not in vocab
|
||
t0, t1 = t[0], t[1]
|
||
for v, vpop in vocab.items():
|
||
if v[0] != t0 and v[0] != t1 and (len(v) < 2 or (v[1] != t0 and v[1] != t1)):
|
||
continue # same first-two-chars filter as search()
|
||
if abs(len(v) - len(t)) > 1 or v == t or not _within1(t, v):
|
||
continue
|
||
phrase = " ".join(toks[:i] + [v] + toks[i + 1:])
|
||
score = (unknown, vpop)
|
||
if cands.get(phrase, (False, -1)) < score:
|
||
cands[phrase] = score
|
||
ranked = sorted(cands, key=cands.get, reverse=True)
|
||
return ranked[:max_phrases]
|