telegram-bot/bot.py

98 lines
3.5 KiB
Python
Raw Normal View History

import logging
import requests
import re
from telegram import InlineQueryResultArticle, InputTextMessageContent
from telegram.ext import Updater, CommandHandler, InlineQueryHandler
token_file = open("token", "r")
updater = Updater(token=token_file.read(), use_context=True)
token_file.close()
dispatcher = updater.dispatcher
logging.basicConfig(format='%(acstime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
class Query:
station_id = -1
departure_count = 4
line = ''
destination = ""
def __init__(self, args):
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 search_station(query):
request = requests.get("https://efa-api.asw.io/api/v1/station/?search=" + query)
if request.status_code != 200 or (request.status_code == 200 and len(request.text) <= 2):
return -1
else:
return request.json()[0]['stationId']
def get_vvs_departures(update, context):
if len(context.args) == 0:
update.message.reply_text("No argument specified")
return
query = Query(context.args)
if query.station_id == -1:
update.message.reply_text("No station matching this name found!")
return
request = requests.get("https://efa-api.asw.io/api/v1/station/" + query.station_id + "/departures")
if request.status_code != 200:
update.message.reply_text("Error with server communication")
update.message.reply_text("Next departures:")
printed_departures = 0
for station in request.json():
if station['direction'].casefold().find(query.destination) != -1 or station['direction'].find(
query.destination) != -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
def inline_station_search(update, context):
query = update.inline_query.query
if len(query) < 5 or not query:
return
results = list()
for station in get_station_id_list(query):
results.append(
InlineQueryResultArticle(
id=station['stationId'],
title=station['name'],
input_message_content=InputTextMessageContent("/vvs " + station['stationId'])
)
)
context.bot.answer_inline_query(update.inline_query.id, results)
def get_station_id_list(name):
request = requests.get("https://efa-api.asw.io/api/v1/station/?search=" + name)
return request.json()
inline_station_search_handler = InlineQueryHandler(inline_station_search)
dispatcher.add_handler(inline_station_search_handler)
updater.dispatcher.add_handler(CommandHandler('vvs', get_vvs_departures))
updater.start_polling()
updater.idle()