telegram-bot/push_information.py

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
from weather_openweathermap import City, get_city_by_name

users = {}
subs_in_edit = {}


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:
            sub.city.query_weather()
            schedule.default_scheduler.every().day.at(time.strftime("%H:%M", sub.time)).do(
                send_message, chat_id=user.chat_id, message=str(sub.city.weather))
            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}!\n"
    string += f"You 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"))
    reply_markup = InlineKeyboardMarkup.from_column(button_list)
    update.message.reply_text(string, reply_markup=reply_markup)
    return "list"


def edit_sub_callback(update: Update, context: CallbackContext):
    global users, subs_in_edit
    user = users[update.effective_chat.id]
    sub_idx = int(update.callback_query.data[11:])
    sub = user.subs[sub_idx]
    subs_in_edit[update.effective_chat.id] = 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}")]]
    reply_markup = InlineKeyboardMarkup(button_list)
    update.effective_chat.send_message("Change subscription.\nWhat do you want to change?", reply_markup=reply_markup)
    return "edit"


def edit_city_callback(update: Update, context: CallbackContext):
    update.effective_chat.send_message("Tell me the city you want information for:")
    return "edit_city"


def edited_city_callback(update: Update, context: CallbackContext):
    global users, subs_in_edit
    user = users[update.effective_chat.id]
    city_name = update.message.text
    city = get_city_by_name(city_name)
    if city.name == "none":
        update.message.reply_text("This city was not found! Try again")
        return None
    else:
        subs_in_edit[user.chat_id].edit_city(city)
        update.message.reply_text(f"City changed to {city.name}.")
        del subs_in_edit[update.effective_chat.id]
        update.message.reply_text(f"Subscription now is {subs_in_edit[user.chat_id]}")
        return ConversationHandler.END


def edit_time_callback(update: Update, context: CallbackContext):
    update.effective_chat.send_message("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 = subs_in_edit[user.chat_id]
    time_string = update.message.text
    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"))
    del subs_in_edit[user.chat_id]
    update.message.reply_text(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.effective_chat.send_message("Deleted")
    schedule_subscriptions()
    return ConversationHandler.END


def add_sub_callback(update: Update, context: CallbackContext):
    update.effective_chat.send_message("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, subs_in_edit
    user = users[update.effective_chat.id]
    city_name = update.message.text
    city = get_city_by_name(city_name)
    if city.name == "none":
        update.message.reply_text("This city was not found! Try again")
        return None
    else:
        update.message.reply_text(f"I found {city.name}. When do you want to receive updates?")
        subs_in_edit[user.chat_id] = Subscription(user.chat_id, city)
    return "time"


def time_callback(update: Update, context: CallbackContext):
    global users, subs_in_edit
    user = users[update.effective_chat.id]
    sub_in_edit = subs_in_edit[update.effective_chat.id]
    time_string = update.message.text
    if time_string.find(':') == -1:
        sub_in_edit.time = time.strptime(time_string, "%H")
    else:
        sub_in_edit.time = time.strptime(time_string, "%H:%M")
    user.add_subscription(sub_in_edit)
    update.message.reply_text(f"Subscription now is: {str(sub_in_edit)}")
    del subs_in_edit[update.effective_chat.id]
    return ConversationHandler.END


def cancel_callback(update: Update, context: CallbackContext):
    update.message.reply_text("Timeout")
    return ConversationHandler.END


def __init__() -> ConversationHandler:
    load_users_from_file()

    push_handler = CommandHandler(command="push", 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)

    time_handler = MessageHandler(callback=time_callback, filters=Filters.text)

    fallback_handler = CommandHandler(command="cancel", callback=cancel_callback)

    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],
                                       "time": [time_handler]},
                               fallbacks=[fallback_handler], conversation_timeout=10)