scripts/update_public_transport_times.py

   1 #!/usr/bin/python
   2 import argparse
   3 import sys
   4 from collections import defaultdict
   5 from copy import deepcopy
   6 from datetime import datetime, date, time, timedelta
   7 from typing import List, Iterable, Optional, Dict, FrozenSet, Final
   8
   9 import jsonschema
  10 import numpy as np
  11 import partridge as ptg
  12 from partridge.gtfs import Feed
  13 from partridge.readers import read_service_ids_by_date
  14 from termcolor import cprint  # python3-termcolor
  15
  16 from wrpylib.cli_tools import unified_diff, input_yes_no_quit, Choice
  17 from wrpylib.json_tools import order_json_keys, format_json
  18 from wrpylib.lib_update_public_transport import default_query_date
  19 from wrpylib.mwapi import WikiSite, page_json
  20 from wrpylib.sledrun_json import Sledrun, PublicTransportStop, PublicTransportStopLine
  21
  22 np.unicode = np.unicode_  # Prevent "AttributeError: module 'numpy' has no attribute 'unicode'"
  23
  24
  25 ROUTE_TYPE: Final = {
  26     0: 'Straßenbahn',
  27     1: 'U-Bahn',
  28     2: 'Bahn',
  29     3: 'Bus',
  30     4: 'Fähre',
  31     5: 'Kabelstraßenbahn',
  32     6: 'Seilbahn',
  33     7: 'Standseilbahn',
  34     11: 'Oberleitungsbus',
  35     12: 'Einschienenbahn',
  36 }
  37
  38
  39 def format_time(day: date, seconds: float) -> str:
  40     dt = datetime.combine(day, time.min)
  41     return (dt + timedelta(seconds=seconds)).isoformat(timespec='minutes')
  42
  43
  44 def update_sledrun(site: WikiSite, title: str, feed: Feed, service_ids_by_date: Dict[datetime.date, FrozenSet[str]]):
  45     cprint(title, 'green')
  46     sledrun_json_page = site.query_page(f'{title}/Rodelbahn.json')
  47     sledrun: Sledrun = page_json(sledrun_json_page)
  48     sledrun_orig = deepcopy(sledrun)
  49
  50     service_date = default_query_date(date.today())
  51
  52     stops = feed.stops
  53     stop_times = feed.stop_times
  54     trips = feed.trips
  55     routes = feed.routes
  56     agency = feed.agency
  57
  58     for pt_stop in sledrun.get('public_transport_stops', []):
  59         pt_stop: PublicTransportStop = pt_stop
  60         ifopt_stop_id = pt_stop.get('ifopt_stop_id')  # e.g. "at:47:61646"
  61         if ifopt_stop_id is None:
  62             continue
  63
  64         selected_stops = stops[stops.stop_id.str.startswith(f'{ifopt_stop_id}:') | (stops.stop_id == ifopt_stop_id)]
  65         if len(selected_stops) == 0:
  66             continue
  67
  68         selected_stop_times = stop_times.merge(selected_stops.stop_id)
  69         selected_trips = trips.merge(selected_stop_times)
  70         selected_trips = selected_trips[selected_trips.service_id.isin(service_ids_by_date[service_date])]
  71         selected_routes = routes.merge(selected_trips.route_id.drop_duplicates())
  72         selected_routes = selected_routes.merge(agency)
  73
  74         selected_trip_stop_times = stop_times.merge(selected_trips.trip_id)
  75         selected_trip_stop_times_grouper = selected_trip_stop_times.groupby('trip_id')['stop_sequence']
  76         first_stops = selected_trip_stop_times.loc[selected_trip_stop_times_grouper.idxmin()]
  77         first_stops = first_stops.merge(stops)
  78         last_stops = selected_trip_stop_times.loc[selected_trip_stop_times_grouper.idxmax()]
  79         last_stops = last_stops.merge(stops)
  80
  81         lines: List[PublicTransportStopLine] = []
  82         for route in selected_routes.itertuples():
  83             departures_dict: Dict[str, List[str]] = defaultdict(list)
  84             arrivals_dict: Dict[str, List[str]] = defaultdict(list)
  85
  86             for trip in selected_trips[selected_trips.route_id == route.route_id].itertuples():
  87                 trip_first_stop = first_stops[first_stops.trip_id == trip.trip_id].iloc[0]
  88                 trip_last_stop = last_stops[last_stops.trip_id == trip.trip_id].iloc[0]
  89                 if trip.stop_id != trip_first_stop.stop_id:
  90                     arrivals_dict[trip_first_stop.stop_name].append(format_time(service_date, trip.arrival_time))
  91                 if trip.stop_id != trip_last_stop.stop_id:
  92                     departures_dict[trip_last_stop.stop_name].append(format_time(service_date, trip.departure_time))
  93
  94             schedule = {
  95                 'service_date': service_date.isoformat(),
  96                 'day_type': 'work_day',
  97                 'departure': [{'direction': k, 'datetime': v} for k, v in departures_dict.items()],
  98                 'arrival': [{'origin': k, 'datetime': v} for k, v in arrivals_dict.items()],
  99             }
 100
 101             lines.append({
 102                 'vao_line_id': f'vvt-{route.route_id}',
 103                 'line': route.route_short_name,
 104                 'category': ROUTE_TYPE[route.route_type],
 105                 'operator': route.agency_name,
 106                 'schedules': [schedule],
 107             })
 108         pt_stop['lines'] = lines
 109
 110     if sledrun == sledrun_orig:
 111         return
 112
 113     jsonschema.validate(instance=sledrun, schema=site.sledrun_schema())
 114     sledrun_ordered = order_json_keys(sledrun, site.sledrun_schema())
 115     assert sledrun_ordered == sledrun
 116     sledrun_orig_str = format_json(sledrun_orig)
 117     sledrun_str = format_json(sledrun_ordered)
 118
 119     unified_diff(sledrun_orig_str, sledrun_str)
 120     choice = input_yes_no_quit('Do you accept the changes [yes, no, quit]? ', None)
 121     if choice == Choice.no:
 122         return
 123     elif choice == Choice.quit:
 124         sys.exit(0)
 125
 126     site(
 127         'edit',
 128         pageid=sledrun_json_page['pageid'],
 129         text=sledrun_str,
 130         summary='Fahrplan zu Haltestellen ergän.',
 131         bot=1,
 132         baserevid=sledrun_json_page['revisions'][0]['revid'],
 133         nocreate=1,
 134         token=site.token(),
 135     )
 136
 137
 138 def get_all_sledrun_titles(site: WikiSite) -> Iterable[str]:
 139     for result in site.query(list='categorymembers', cmtitle='Kategorie:Rodelbahn', cmlimit='max'):
 140         for page in result['categorymembers']:
 141             yield page['title']
 142
 143 def update_public_transport_times(ini_files: List[str], gtfs_file: str, sledrun_title: Optional[str]):
 144     feed = ptg.load_feed(gtfs_file)
 145     service_ids_by_date = read_service_ids_by_date(gtfs_file)
 146
 147     site = WikiSite(ini_files)
 148     if sledrun_title is None:
 149         for sledrun_title in get_all_sledrun_titles(site):
 150             update_sledrun(site, sledrun_title, feed, service_ids_by_date)
 151     else:
 152         update_sledrun(site, sledrun_title, feed, service_ids_by_date)
 153
 154
 155 def main():
 156     parser = argparse.ArgumentParser(description='Update public transport bus stop time info in sledrun JSON files.')
 157     parser.add_argument('--sledrun', help='If given, work on a single sled run page, otherwise at the whole category.')
 158     parser.add_argument('gtfs_file', help='GTFS file.')
 159     parser.add_argument('inifile', nargs='+', help='inifile.ini, see: https://www.winterrodeln.org/trac/wiki/ConfigIni')
 160     args = parser.parse_args()
 161     update_public_transport_times(args.inifile, args.gtfs_file, args.sledrun)
 162
 163
 164 if __name__ == '__main__':
 165     main()