#!/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, service_date: date, 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) 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': sorted(v)} for k, v in departures_dict.items()], 'arrival': [{'origin': k, 'datetime': sorted(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änzt.', 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], service_date: date, 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, service_date, feed, service_ids_by_date) else: update_sledrun(site, sledrun_title, service_date, feed, service_ids_by_date) def main(): query_date = default_query_date(date.today()) 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('--date', type=date.fromisoformat, default=query_date, help='Working week date to query the database.') 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.date, args.gtfs_file, args.sledrun) if __name__ == '__main__': main()