Add script to update public transport times.

diff --git a/scripts/update_public_transport_times.py b/scripts/update_public_transport_times.py
new file mode 100644 (file)
index 0000000..c97650c
--- /dev/null
+++ b/scripts/update_public_transport_times.py
@@ -0,0 +1,165 @@
+#!/usr/bin/python
+import argparse
+import sys
+from collections import defaultdict
+from copy import deepcopy
+from datetime import datetime, date, time, timedelta
+from typing import List, Iterable, Optional, Dict, FrozenSet, Final
+
+import jsonschema
+import numpy as np
+import partridge as ptg
+from partridge.gtfs import Feed
+from partridge.readers import read_service_ids_by_date
+from termcolor import cprint  # python3-termcolor
+
+from wrpylib.cli_tools import unified_diff, input_yes_no_quit, Choice
+from wrpylib.json_tools import order_json_keys, format_json
+from wrpylib.lib_update_public_transport import default_query_date
+from wrpylib.mwapi import WikiSite, page_json
+from wrpylib.sledrun_json import Sledrun, PublicTransportStop, PublicTransportStopLine
+
+np.unicode = np.unicode_  # Prevent "AttributeError: module 'numpy' has no attribute 'unicode'"
+
+
+ROUTE_TYPE: Final = {
+    0: 'Straßenbahn',
+    1: 'U-Bahn',
+    2: 'Bahn',
+    3: 'Bus',
+    4: 'Fähre',
+    5: 'Kabelstraßenbahn',
+    6: 'Seilbahn',
+    7: 'Standseilbahn',
+    11: 'Oberleitungsbus',
+    12: 'Einschienenbahn',
+}
+
+
+def format_time(day: date, seconds: float) -> str:
+    dt = datetime.combine(day, time.min)
+    return (dt + timedelta(seconds=seconds)).isoformat(timespec='minutes')
+
+
+def update_sledrun(site: WikiSite, title: str, feed: Feed, service_ids_by_date: Dict[datetime.date, FrozenSet[str]]):
+    cprint(title, 'green')
+    sledrun_json_page = site.query_page(f'{title}/Rodelbahn.json')
+    sledrun: Sledrun = page_json(sledrun_json_page)
+    sledrun_orig = deepcopy(sledrun)
+
+    service_date = default_query_date(date.today())
+
+    stops = feed.stops
+    stop_times = feed.stop_times
+    trips = feed.trips
+    routes = feed.routes
+    agency = feed.agency
+
+    for pt_stop in sledrun.get('public_transport_stops', []):
+        pt_stop: PublicTransportStop = pt_stop
+        ifopt_stop_id = pt_stop.get('ifopt_stop_id')  # e.g. "at:47:61646"
+        if ifopt_stop_id is None:
+            continue
+
+        selected_stops = stops[stops.stop_id.str.startswith(f'{ifopt_stop_id}:') | (stops.stop_id == ifopt_stop_id)]
+        if len(selected_stops) == 0:
+            continue
+
+        selected_stop_times = stop_times.merge(selected_stops.stop_id)
+        selected_trips = trips.merge(selected_stop_times)
+        selected_trips = selected_trips[selected_trips.service_id.isin(service_ids_by_date[service_date])]
+        selected_routes = routes.merge(selected_trips.route_id.drop_duplicates())
+        selected_routes = selected_routes.merge(agency)
+
+        selected_trip_stop_times = stop_times.merge(selected_trips.trip_id)
+        selected_trip_stop_times_grouper = selected_trip_stop_times.groupby('trip_id')['stop_sequence']
+        first_stops = selected_trip_stop_times.loc[selected_trip_stop_times_grouper.idxmin()]
+        first_stops = first_stops.merge(stops)
+        last_stops = selected_trip_stop_times.loc[selected_trip_stop_times_grouper.idxmax()]
+        last_stops = last_stops.merge(stops)
+
+        lines: List[PublicTransportStopLine] = []
+        for route in selected_routes.itertuples():
+            departures_dict: Dict[str, List[str]] = defaultdict(list)
+            arrivals_dict: Dict[str, List[str]] = defaultdict(list)
+
+            for trip in selected_trips[selected_trips.route_id == route.route_id].itertuples():
+                trip_first_stop = first_stops[first_stops.trip_id == trip.trip_id].iloc[0]
+                trip_last_stop = last_stops[last_stops.trip_id == trip.trip_id].iloc[0]
+                if trip.stop_id != trip_first_stop.stop_id:
+                    arrivals_dict[trip_first_stop.stop_name].append(format_time(service_date, trip.arrival_time))
+                if trip.stop_id != trip_last_stop.stop_id:
+                    departures_dict[trip_last_stop.stop_name].append(format_time(service_date, trip.departure_time))
+
+            schedule = {
+                'service_date': service_date.isoformat(),
+                'day_type': 'work_day',
+                'departure': [{'direction': k, 'datetime': v} for k, v in departures_dict.items()],
+                'arrival': [{'origin': k, 'datetime': v} for k, v in arrivals_dict.items()],
+            }
+
+            lines.append({
+                'vao_line_id': f'vvt-{route.route_id}',
+                'line': route.route_short_name,
+                'category': ROUTE_TYPE[route.route_type],
+                'operator': route.agency_name,
+                'schedules': [schedule],
+            })
+        pt_stop['lines'] = lines
+
+    if sledrun == sledrun_orig:
+        return
+
+    jsonschema.validate(instance=sledrun, schema=site.sledrun_schema())
+    sledrun_ordered = order_json_keys(sledrun, site.sledrun_schema())
+    assert sledrun_ordered == sledrun
+    sledrun_orig_str = format_json(sledrun_orig)
+    sledrun_str = format_json(sledrun_ordered)
+
+    unified_diff(sledrun_orig_str, sledrun_str)
+    choice = input_yes_no_quit('Do you accept the changes [yes, no, quit]? ', None)
+    if choice == Choice.no:
+        return
+    elif choice == Choice.quit:
+        sys.exit(0)
+
+    site(
+        'edit',
+        pageid=sledrun_json_page['pageid'],
+        text=sledrun_str,
+        summary='Fahrplan zu Haltestellen ergän.',
+        bot=1,
+        baserevid=sledrun_json_page['revisions'][0]['revid'],
+        nocreate=1,
+        token=site.token(),
+    )
+
+
+def get_all_sledrun_titles(site: WikiSite) -> Iterable[str]:
+    for result in site.query(list='categorymembers', cmtitle='Kategorie:Rodelbahn', cmlimit='max'):
+        for page in result['categorymembers']:
+            yield page['title']
+
+def update_public_transport_times(ini_files: List[str], gtfs_file: str, sledrun_title: Optional[str]):
+    feed = ptg.load_feed(gtfs_file)
+    service_ids_by_date = read_service_ids_by_date(gtfs_file)
+
+    site = WikiSite(ini_files)
+    if sledrun_title is None:
+        for sledrun_title in get_all_sledrun_titles(site):
+            update_sledrun(site, sledrun_title, feed, service_ids_by_date)
+    else:
+        update_sledrun(site, sledrun_title, feed, service_ids_by_date)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Update public transport bus stop time info in sledrun JSON files.')
+    parser.add_argument('--sledrun', help='If given, work on a single sled run page, otherwise at the whole category.')
+    parser.add_argument('gtfs_file', help='GTFS file.')
+    parser.add_argument('inifile', nargs='+', help='inifile.ini, see: https://www.winterrodeln.org/trac/wiki/ConfigIni')
+    args = parser.parse_args()
+    update_public_transport_times(args.inifile, args.gtfs_file, args.sledrun)
+
+
+if __name__ == '__main__':
+    main()