diff --git a/bot.py b/bot.py index 6d98cae..d0ee8c4 100644 --- a/bot.py +++ b/bot.py @@ -1,5 +1,6 @@ import logging import requests +import re from telegram import InlineQueryResultArticle, InputTextMessageContent from telegram.ext import Updater, CommandHandler, InlineQueryHandler @@ -13,28 +14,22 @@ logging.basicConfig(format='%(acstime)s - %(name)s - %(levelname)s - %(message)s class Query: - input = "" station_id = -1 departure_count = 4 - line_filter = -1 + line = '' destination = "" def __init__(self, args): - last_index_string = len(args) - if args[-1].isdigit(): - self.departure_count = int(args[-1]) - last_index_string -= 1 - request_tmp = ' '.join(args[0:last_index_string]).strip() - separator_index = request_tmp.find(" to ") - if separator_index != -1: - self.input = request_tmp[:separator_index - 1].strip() - self.destination = request_tmp[separator_index + 4:].strip() - else: - self.input = request_tmp.strip() - if args[0].isdigit(): - self.station_id = args[0].strip() - else: - self.station_id = search_station(self.input) + request_tmp = ' '.join(args) + argument_names = re.findall(r' to | in | times ', request_tmp) + arguments = re.split(r' to | in | times ', request_tmp) + self.station_id = search_station(arguments[0]) + if ' to ' in argument_names: + self.destination = arguments[argument_names.index(' to ')+1] + if ' in ' in argument_names: + self.line = arguments[argument_names.index(' in ')+1] + if ' times ' in argument_names: + self.departure_count = int(arguments[argument_names.index(' times ') + 1]) def start(update): @@ -65,11 +60,12 @@ def get_vvs_departures(update, context): for station in request.json(): if station['direction'].casefold().find(query.destination) != -1 or station['direction'].find( query.destination) != -1: - update.message.reply_text( - "Line " + station['number'] + " to \"" + station['direction'] + "\" at " - + station['departureTime']['hour'].zfill(2) + ':' + station['departureTime']['minute'].zfill(2) - + " +" + str(station['delay'])) - printed_departures += 1 + if query.line == '' or (query.line != -1 and station['number'] == query.line): + update.message.reply_text( + "Line " + station['number'] + " to \"" + station['direction'] + "\" at " + + station['departureTime']['hour'].zfill(2) + ':' + station['departureTime']['minute'].zfill(2) + + " +" + str(station['delay'])) + printed_departures += 1 if printed_departures >= query.departure_count: break