135 lines
4.1 KiB
Python
135 lines
4.1 KiB
Python
"""Load the checked-in Austrian stop snapshot into the ``stops`` table.
|
|
|
|
The snapshot (data/stops_at.csv.gz) is produced by scripts/fetch_stops.py from
|
|
OpenStreetMap data (© OpenStreetMap contributors, ODbL). Importing needs no
|
|
network access.
|
|
|
|
flask import-stops # import the checked-in snapshot
|
|
flask import-stops --file X # import an alternative CSV(.gz)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import gzip
|
|
import io
|
|
import unicodedata
|
|
from pathlib import Path
|
|
|
|
import click
|
|
from flask import current_app
|
|
from flask.cli import with_appcontext
|
|
|
|
from models import Stop, Trip, db
|
|
|
|
SNAPSHOT_PATH = Path(__file__).resolve().parent / "data" / "stops_at.csv.gz"
|
|
|
|
|
|
def normalize_name(value: str) -> str:
|
|
"""Lowercase, fold accents and ß for diacritic-insensitive search.
|
|
|
|
Mirrors scripts/fetch_stops.py:normalize_name - keep the two in sync.
|
|
"""
|
|
value = value.strip().lower().replace("ß", "ss")
|
|
decomposed = unicodedata.normalize("NFKD", value)
|
|
return "".join(ch for ch in decomposed if not unicodedata.combining(ch))
|
|
|
|
|
|
def _open_csv(path: Path) -> io.TextIOBase:
|
|
if str(path).endswith(".gz"):
|
|
return gzip.open(path, "rt", newline="", encoding="utf-8")
|
|
return open(path, "rt", newline="", encoding="utf-8")
|
|
|
|
|
|
def import_stops(path: Path | str = SNAPSHOT_PATH) -> dict[str, int]:
|
|
"""Upsert every row of the snapshot into ``stops`` by (osm_type, osm_id).
|
|
|
|
Stops that vanished from the snapshot are removed unless a trip still
|
|
references them, so foreign keys from ``trips`` stay valid.
|
|
"""
|
|
path = Path(path)
|
|
if not path.exists():
|
|
raise click.ClickException(
|
|
f"Snapshot not found: {path}\nRun scripts/fetch_stops.py first."
|
|
)
|
|
|
|
with _open_csv(path) as fh:
|
|
rows = list(csv.DictReader(fh))
|
|
if not rows:
|
|
raise click.ClickException(f"{path} contains no rows.")
|
|
|
|
existing = {(s.osm_type, s.osm_id): s for s in Stop.query.all()}
|
|
seen: set[tuple[str, int]] = set()
|
|
inserted = updated = 0
|
|
|
|
for row in rows:
|
|
key = (row["osm_type"], int(row["osm_id"]))
|
|
seen.add(key)
|
|
name = row["name"].strip()
|
|
fields = dict(
|
|
name=name,
|
|
name_normalized=normalize_name(name),
|
|
stop_type=row["stop_type"] or "other",
|
|
municipality=(row.get("municipality") or "").strip() or None,
|
|
latitude=float(row["latitude"]),
|
|
longitude=float(row["longitude"]),
|
|
)
|
|
stop = existing.get(key)
|
|
if stop is None:
|
|
db.session.add(Stop(osm_type=key[0], osm_id=key[1], **fields))
|
|
inserted += 1
|
|
else:
|
|
changed = False
|
|
for attr, value in fields.items():
|
|
if getattr(stop, attr) != value:
|
|
setattr(stop, attr, value)
|
|
changed = True
|
|
updated += changed
|
|
|
|
referenced = {
|
|
sid
|
|
for (sid,) in db.session.query(Trip.origin_stop_id).distinct()
|
|
if sid is not None
|
|
} | {
|
|
sid
|
|
for (sid,) in db.session.query(Trip.destination_stop_id).distinct()
|
|
if sid is not None
|
|
}
|
|
removed = 0
|
|
for key, stop in existing.items():
|
|
if key not in seen and stop.id not in referenced:
|
|
db.session.delete(stop)
|
|
removed += 1
|
|
|
|
db.session.commit()
|
|
return {
|
|
"total": len(rows),
|
|
"inserted": inserted,
|
|
"updated": updated,
|
|
"removed": removed,
|
|
}
|
|
|
|
|
|
@click.command("import-stops")
|
|
@click.option(
|
|
"--file",
|
|
"file_path",
|
|
type=click.Path(exists=True, dir_okay=False),
|
|
default=None,
|
|
help="CSV(.gz) to import instead of the checked-in snapshot.",
|
|
)
|
|
@with_appcontext
|
|
def import_stops_command(file_path: str | None) -> None:
|
|
"""Import Austrian public-transport stops into the database."""
|
|
stats = import_stops(file_path or SNAPSHOT_PATH)
|
|
click.echo(
|
|
"Stops imported: "
|
|
f"{stats['total']} in snapshot, "
|
|
f"{stats['inserted']} new, {stats['updated']} updated, "
|
|
f"{stats['removed']} removed. "
|
|
f"Total now: {Stop.query.count()}."
|
|
)
|
|
|
|
|
|
def register_cli(app) -> None:
|
|
app.cli.add_command(import_stops_command)
|