import re from datetime import timezone, timedelta import requests from dateutil.parser import parse from exceptions import ServerCommunicationError, StationNotFoundError parse_regex = r' to | times | in ' def handle_marudor_departures(update, context): command = ' '.join(context.args) arguments = parse_arguments(command) eva_id = get_eva_id_for_station_name(arguments['station']) departures = get_departures(eva_id) if ' in ' in arguments: departures = filter_line_names(departures, arguments[' in ']) if ' to ' in arguments: departures = filter_destination(departures, arguments[' to ']) departure_count = 4 if ' times ' in arguments: departure_count = int(arguments[' times ']) departures = departures[:departure_count] for reply in build_reply(departures): update.message.reply_text(reply) def parse_arguments(command: str) -> dict: included_commands = re.findall(parse_regex, command) parameters = re.split(parse_regex, command) return dict(zip(['station'] + included_commands, parameters)) def get_eva_id_for_station_name(station_name: str) -> str: request = requests.get(f'https://marudor.de/api/hafas/v1/station/{station_name}') if request.status_code != 200: raise ServerCommunicationError response = request.json() if len(response) == 0 or 'id' not in response[0]: raise StationNotFoundError return response[0]['id'] def get_departures(eva_id: str) -> list: request = requests.get(f'https://marudor.de/api/hafas/experimental/irisCompatibleAbfahrten/{eva_id}?lookbehind=1') if request.status_code != 200: raise ServerCommunicationError response = request.json() return response['departures'] def filter_destination(departures: list, destination: str) -> list: departures[:] = [departure for departure in departures if destination.casefold() in departure['destination'].casefold() or destination.casefold() in departure[ 'scheduledDestination'].casefold()] return departures def filter_line_names(departures: list, line_name: str) -> list: departures[:] = [departure for departure in departures if line_name in departure['train']['name']] return departures def build_reply(departures: list) -> list: if len(departures) == 0: return ['No departures found!'] reply = [f'Next departures for station: {departures[0]["currentStation"]["title"]}'] for departure in departures: reply_text = f'{departure["train"]["name"]} to "{departure["destination"]}" at ' \ f'{parse(departure["departure"]["time"]).astimezone(timezone(timedelta(hours=1))).strftime("%H:%M")}' if 'delay' in departure['departure']: reply_text += f' (+{str(departure["departure"]["delay"])})' reply.append(reply_text) return reply