telegram-bot/push_information.py
2020-09-25 16:24:22 +02:00

286 lines
12 KiB
Python

import _pickle as pickle
import time
import schedule
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import ConversationHandler, CommandHandler, MessageHandler, CallbackQueryHandler, CallbackContext
from telegram.ext.filters import Filters
from bot import send_message, edit_message_and_delete_answer
from weather_openweathermap import City, get_city_by_name
users = {}
LIST, ADD_CITY, ADD_TIME, EDIT, EDIT_CITY, EDIT_TIME = range(6)
class Subscription(object):
time: time
city: City
chat_id: int
def __init__(self, chat_id: int, city: City):
self.chat_id = chat_id
self.city = city
def __str__(self):
return f"{self.city.name} at {time.strftime('%H:%M', self.time)}"
def edit_time(self, edited_time: time):
self.time = edited_time
schedule_subscriptions()
save_users_to_file()
def edit_city(self, edited_city: City):
self.city = edited_city
schedule_subscriptions()
save_users_to_file()
class User(object):
chat_id: int
subs = []
def __init__(self, chat_id: int):
self.chat_id = chat_id
global users
users[self.chat_id] = self.subs
def __getstate__(self):
attributes = self.__dict__.copy()
attributes['subs'] = self.subs
for i in range(len(attributes['subs'])):
try:
del attributes['subs'][i].city.weather
except AttributeError:
pass
return attributes
def __setstate__(self, state):
self.chat_id = state['chat_id']
self.subs = state['subs']
def add_subscription(self, subscription: Subscription):
self.subs.append(subscription)
global users
users[self.chat_id] = self
save_users_to_file()
schedule_subscriptions()
def del_subscription(self, idx: int):
del self.subs[idx]
save_users_to_file()
schedule_subscriptions()
def save_users_to_file():
global users
with open("users.pkl", 'wb') as users_file:
pickle.dump(users, users_file)
def load_users_from_file():
global users
try:
with open("users.pkl", 'rb') as users_file:
users = pickle.load(users_file)
schedule_subscriptions()
except FileNotFoundError:
pass
except EOFError:
pass
def schedule_subscriptions():
schedule.default_scheduler.clear()
for user in users.values():
for sub in user.subs:
schedule.default_scheduler.every().day.at(time.strftime("%H:%M", sub.time)).do(handle_subscription,
chat_id=user.chat_id,
sub=sub)
def handle_subscription(chat_id: int, sub: Subscription):
sub.city.query_weather()
send_message(chat_id=chat_id, message=sub.city.get_weather_str())
sub.city.weather.clear()
def push_callback(update: Update, context: CallbackContext):
global users
if update.effective_chat.id in users:
user = users[update.effective_chat.id]
else:
user = User(update.effective_chat.id)
users[update.effective_chat.id] = user
string = f"Hello, {update.effective_user.first_name}!\nYou have {len(user.subs)} subscriptions:\n"
button_list = []
for i, subscription in enumerate(user.subs):
button_list.append(
InlineKeyboardButton(text=f"{subscription.city.name} at {time.strftime('%H:%M', subscription.time)}",
callback_data=f"/push_edit {i}"))
button_list.append(InlineKeyboardButton(text="Add", callback_data="/push_add"))
button_list.append(InlineKeyboardButton(text="End", callback_data="/cancel"))
reply_markup = InlineKeyboardMarkup.from_column(button_list)
update.message.reply_text(string, reply_markup=reply_markup)
context.bot.delete_message(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id)
return LIST
def edit_sub_callback(update: Update, context: CallbackContext):
global users
user = users[update.effective_chat.id]
sub_idx = int(update.callback_query.data[11:])
sub = user.subs[sub_idx]
context.user_data['sub'] = sub
button_list = [[InlineKeyboardButton(text="City", callback_data="/push_edit_city"),
InlineKeyboardButton(text="Time", callback_data="/push_edit_time")],
[InlineKeyboardButton(text="Delete", callback_data=f"/push_edit_delete {sub_idx}"),
InlineKeyboardButton(text="Cancel", callback_data="/cancel")]]
reply_markup = InlineKeyboardMarkup(button_list)
update.callback_query.answer()
update.callback_query.edit_message_text(f"Change subscription {sub}.\nWhat do you want to change?",
reply_markup=reply_markup)
return EDIT
def edit_city_callback(update: Update, context: CallbackContext):
update.callback_query.answer()
update.callback_query.edit_message_text("Tell me the city you want information for:")
return EDIT_CITY
def edited_city_callback(update: Update, context: CallbackContext):
global users
user = users[update.effective_chat.id]
city_name = update.message.text
city = get_city_by_name(city_name)
if city.name == "none":
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message=f"The city {city_name} was not found! Try again")
return None
else:
context.user_data['sub'].edit_city(city)
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message=f"Subscription now is {context.user_data['sub']}")
del context.user_data['sub']
return ConversationHandler.END
def edit_time_callback(update: Update, context: CallbackContext):
update.callback_query.answer()
update.callback_query.edit_message_text("When do you want to receive updates?")
return EDIT_TIME
def edited_time_callback(update: Update, context: CallbackContext):
global users
user = users[update.effective_chat.id]
sub_in_edit = context.user_data['sub']
time_string = update.message.text
try:
if time_string.find(':') == -1:
sub_in_edit.edit_time(time.strptime(time_string, "%H"))
else:
sub_in_edit.edit_time(time.strptime(time_string, "%H:%M"))
except ValueError:
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message="This time was not recognized! Try again")
return None
del context.user_data['sub']
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message=f"Subscription now is {sub_in_edit}")
return ConversationHandler.END
def delete_subscription_callback(update: Update, context: CallbackContext):
global users
users[update.effective_chat.id].del_subscription(int(update.callback_query.data[18:]))
update.callback_query.answer()
update.callback_query.edit_message_text("Deleted")
schedule_subscriptions()
return ConversationHandler.END
def add_sub_callback(update: Update, context: CallbackContext):
update.callback_query.answer()
update.callback_query.edit_message_text("Add a new subscription.\nTell me the city you want information for:")
return ADD_CITY
def add_city_callback(update: Update, context: CallbackContext):
global users
user = users[update.effective_chat.id]
city_name = update.message.text
city = get_city_by_name(city_name)
if city.name == "none":
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message="This city was not found! Try again")
return None
else:
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message=f"I found {city.name}. When do you want to receive updates?")
context.user_data['sub'] = Subscription(user.chat_id, city)
return ADD_TIME
def add_time_callback(update: Update, context: CallbackContext):
global users
user = users[update.effective_chat.id]
sub_in_edit = context.user_data['sub']
time_string = update.message.text
try:
if time_string.find(':') == -1:
sub_in_edit.edit_time(time.strptime(time_string, "%H"))
else:
sub_in_edit.edit_time(time.strptime(time_string, "%H:%M"))
except ValueError:
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message="This time was not recognized! Try again")
return None
user.add_subscription(sub_in_edit)
edit_message_and_delete_answer(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id,
message=f"Subscription now is: {str(sub_in_edit)}")
del context.user_data['sub']
return ConversationHandler.END
def cancel_callback(update: Update, context: CallbackContext):
context.bot.delete_message(chat_id=update.effective_chat.id, message_id=update.effective_message.message_id)
return ConversationHandler.END
def cancel_button_callback(update: Update, context: CallbackContext):
update.callback_query.answer()
update.effective_message.delete()
return ConversationHandler.END
def __init__() -> ConversationHandler:
load_users_from_file()
push_handler = CommandHandler(command=["push", "subs"], callback=push_callback)
edit_sub_handler = CallbackQueryHandler(callback=edit_sub_callback, pattern="^\/push_edit ")
edit_city_handler = CallbackQueryHandler(callback=edit_city_callback, pattern="^\/push_edit_city")
edited_city_handler = MessageHandler(callback=edited_city_callback, filters=Filters.text)
edit_time_handler = CallbackQueryHandler(callback=edit_time_callback, pattern="^\/push_edit_time")
edited_time_handler = MessageHandler(callback=edited_time_callback, filters=Filters.text)
delete_subscription_handler = CallbackQueryHandler(callback=delete_subscription_callback,
pattern="^\/push_edit_delete")
add_sub_handler = CallbackQueryHandler(callback=add_sub_callback, pattern="^\/push_add")
add_city_handler = MessageHandler(callback=add_city_callback, filters=Filters.text)
add_time_handler = MessageHandler(callback=add_time_callback, filters=Filters.text)
fallback_handler = CommandHandler(command="cancel", callback=cancel_callback)
fallback_callback_handler = CallbackQueryHandler(callback=cancel_button_callback, pattern="^\/cancel")
return ConversationHandler(entry_points=[push_handler],
states={LIST: [edit_sub_handler, add_sub_handler], ADD_CITY: [add_city_handler],
EDIT: [edit_city_handler, edit_time_handler, delete_subscription_handler],
EDIT_CITY: [edited_city_handler], EDIT_TIME: [edited_time_handler],
ADD_TIME: [add_time_handler]},
fallbacks=[fallback_handler, fallback_callback_handler])