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])