From 06bd4dc075b1d2901d4503ba1fcf232a238eccc4 Mon Sep 17 00:00:00 2001 From: Daniel Friesel Date: Sun, 28 Mar 2021 18:58:37 +0200 Subject: initial commit --- bin/lookup-server | 203 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) create mode 100755 bin/lookup-server (limited to 'bin/lookup-server') diff --git a/bin/lookup-server b/bin/lookup-server new file mode 100755 index 0000000..13bff59 --- /dev/null +++ b/bin/lookup-server @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# vim:tabstop=4 softtabstop=4 shiftwidth=4 textwidth=160 smarttab expandtab colorcolumn=160 + +import argparse +import psycopg2 +import aiohttp +from aiohttp import web +from datetime import datetime, timedelta +import dateutil.parser +from geopy.distance import distance +import json +import os +import pytz + +headers = { + "Access-Control-Allow-Origin": "*", + "Content-Type": "application/json; charset=utf-8", +} + +conn = psycopg2.connect( + dbname=os.getenv("GEOLOOKUP_DBNAME", "geo_to_stations"), + user=os.getenv("GEOLOOKUP_DBUSER", "geo_to_stations"), + password=os.getenv("GEOLOOKUP_DBPASS"), + host=os.getenv("GEOLOOKUP_DBHOST", "localhost"), +) + + +def set_coarse_location(train): + now = datetime.now(pytz.utc) + train_evas = None + stopovers = train["previousStopovers"] + + for i, stopover in enumerate(stopovers): + if stopover["departure"]: + stopover["departure"] = dateutil.parser.parse(stopover["departure"]) + + # start with origin. (planned)arrival is always null in a previousStopovers list + departure = stopover["departure"] + if i > 0 and departure and departure > now: + train_evas = ( + int(stopovers[i - 1]["stop"]["id"]), + int(stopover["stop"]["id"]), + ) + train_stops = (stopovers[i - 1]["stop"]["name"], stopover["stop"]["name"]) + train_coords = ( + ( + stopovers[i - 1]["stop"]["location"]["latitude"], + stopovers[i - 1]["stop"]["location"]["longitude"], + ), + ( + stopover["stop"]["location"]["latitude"], + stopover["stop"]["location"]["longitude"], + ), + ) + train_times = (stopovers[i - 1]["departure"], departure) + break + if not train_evas: + train_evas = (int(train["stop"]["id"]), int(stopovers[-1]["stop"]["id"])) + train_stops = (train["stop"]["name"], stopovers[-1]["stop"]["name"]) + train_coords = ( + ( + stopovers[-1]["stop"]["location"]["latitude"], + stopovers[-1]["stop"]["location"]["longitude"], + ), + ( + train["stop"]["location"]["latitude"], + train["stop"]["location"]["longitude"], + ), + ) + train_times = (stopovers[-1]["departure"], dateutil.parser.parse(train["when"])) + + if not train_times[0]: + return + + train["evas"] = train_evas + train["stop_names"] = train_stops + train["coords"] = train_coords + train["times"] = train_times + + train["progress_ratio"] = 1 - ( + (train["times"][1].timestamp() - now.timestamp()) + / (train["times"][1].timestamp() - train["times"][0].timestamp()) + ) + train["progress_ratio"] = max(0, min(1, train["progress_ratio"])) + + if train["progress_ratio"] == 0: + train["location"] = train["coarse_location"] = train["coords"][0] + elif train["progress_ratio"] == 1: + train["location"] = train["coarse_location"] = train["coords"][1] + else: + ratio = train["progress_ratio"] + coords = train["coords"] + train["coarse_location"] = ( + coords[1][0] * ratio + coords[0][0] * (1 - ratio), + coords[1][1] * ratio + coords[0][1] * (1 - ratio), + ) + + +def calculate_distance(train, latlon): + train["distance"] = distance(train["coarse_location"], latlon).km + + +def format_train(train): + return { + "line": train["line"]["name"], + "no": train["line"]["fahrtNr"], + "tripId": train["tripId"], + "location": train["coarse_location"], + "distance": train["distance"], + "stops": [ + (train["evas"][0], train["stop_names"][0], train["times"][0].isoformat()), + (train["evas"][1], train["stop_names"][1], train["times"][1].isoformat()), + ], + } + + +async def handle_search(request): + try: + lat = float(request.query.get("lat")) + lon = float(request.query.get("lon")) + except TypeError: + return web.HTTPBadRequest(text="lat/lon are mandatory") + except ValueError: + return web.HTTPBadRequest(text="lat/lon must be floating-point numbers") + + lut_lat = round(lat * 1000) + lut_lon = round(lon * 1000) + + evas = set() + + with conn.cursor() as cur: + cur.execute( + "select stations from stations where lat between %s and %s and lon between %s and %s", + (lut_lat - 3, lut_lat + 3, lut_lon - 3, lut_lon + 3), + ) + for eva_list in cur.fetchall(): + evas.update(eva_list[0]) + + if not evas: + response = {"evas": list(), "trains": list()} + return web.Response(body=json.dumps(response), headers=headers) + + arrivals = list() + trains = list() + + # deliberately not parallelized to minimize load on transport.rest + for eva in evas: + async with aiohttp.ClientSession() as session: + async with session.get( + f"https://v5.db.transport.rest/stops/{eva}/arrivals?results=40&duration=120&stopovers=true&bus=false&subway=false&tram=false" + ) as response: + content = await response.text() + content = json.loads(content) + arrivals.append(content) + + for train_list in arrivals: + for train in train_list: + is_candidate = False + for stop in train["previousStopovers"]: + if ( + int(stop["stop"]["id"]) in evas + and stop["stop"]["id"] != train["stop"]["id"] + ): + is_candidate = True + break + if is_candidate: + trains.append(train) + + seen = set() + trains = [ + seen.add(train["line"]["fahrtNr"]) or train + for train in trains + if train["line"]["fahrtNr"] not in seen + ] + + for train in trains: + set_coarse_location(train) + + trains = list(filter(lambda train: "coarse_location" in train, trains)) + + for train in trains: + calculate_distance(train, (lat, lon)) + + trains = sorted(trains, key=lambda train: train["distance"]) + trains = list(map(format_train, trains)) + + response = {"evas": list(evas), "trains": trains} + + return web.Response(body=json.dumps(response, ensure_ascii=False), headers=headers) + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser( + description="geolocation to train estimation service" + ) + parser.add_argument("--port", type=int, metavar="PORT", default=8080) + parser.add_argument("--prefix", type=str, metavar="PATH", default="/") + args = parser.parse_args() + + app = web.Application() + app.add_routes([web.get(f"{args.prefix}search", handle_search)]) + web.run_app(app, host="localhost", port=args.port) -- cgit v1.2.3