diff --git a/examples/step_example.py b/examples/step_example.py index fd8d07d..0fc0ef1 100644 --- a/examples/step_example.py +++ b/examples/step_example.py @@ -75,4 +75,14 @@ def process_sex_step(message): bot.reply_to(message, 'oooops') +# Enable saving next step handlers to file "./.handlers-saves/step.save". +# Delay=2 means that after any change in next step handlers (e.g. calling register_next_step_handler()) +# saving will hapen after delay 2 seconds. +bot.enable_save_next_step_handlers(delay=2) + +# Load next_step_handlers from save file (default "./.handlers-saves/step.save") +# WARNING It will work only if enable_save_next_step_handlers was called! +bot.load_next_step_handlers() + + bot.polling() diff --git a/telebot/__init__.py b/telebot/__init__.py index fb7770b..6e4ef40 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -7,6 +7,9 @@ import re import sys import six +import os +import pickle + import logging logger = logging.getLogger('TeleBot') @@ -27,6 +30,70 @@ Module : telebot """ +class Handler: + """ + Class for (next step|reply) handlers + """ + def __init__(self, callback, *args, **kwargs): + self.callback = callback + self.args = args + self.kwargs = kwargs + + def __getitem__(self, item): + return getattr(self, item) + + +class Saver: + """ + Class for saving (next step|reply) handlers + """ + def __init__(self, handlers, filename, delay): + self.handlers = handlers + self.filename = filename + self.delay = delay + self.timer = threading.Timer(delay, self.save_handlers) + + def start_save_timer(self): + if not self.timer.is_alive(): + if self.delay <= 0: + self.save_handlers() + else: + self.timer = threading.Timer(self.delay, self.save_handlers) + self.timer.start() + + def save_handlers(self): + self.dump_handlers(self.handlers, self.filename) + + def load_handlers(self, filename, del_file_after_loading=True): + tmp = self.return_load_handlers(filename, del_file_after_loading=del_file_after_loading) + if tmp is not None: + self.handlers.update(tmp) + + @staticmethod + def dump_handlers(handlers, filename, file_mode="wb"): + dirs = filename.rsplit('/', maxsplit=1)[0] + os.makedirs(dirs, exist_ok=True) + + with open(filename + ".tmp", file_mode) as file: + pickle.dump(handlers, file) + + if os.path.isfile(filename): + os.remove(filename) + + os.rename(filename + ".tmp", filename) + + @staticmethod + def return_load_handlers(filename, del_file_after_loading=True): + if os.path.isfile(filename) and os.path.getsize(filename) > 0: + with open(filename, "rb") as file: + handlers = pickle.load(file) + + if del_file_after_loading: + os.remove(filename) + + return handlers + + class TeleBot: """ This is TeleBot Class Methods: @@ -85,6 +152,9 @@ class TeleBot: # key: chat_id, value: handler list self.next_step_handlers = {} + self.next_step_saver = None + self.reply_saver = None + self.message_handlers = [] self.edited_message_handlers = [] self.channel_post_handlers = [] @@ -99,6 +169,54 @@ class TeleBot: if self.threaded: self.worker_pool = util.ThreadPool(num_threads=num_threads) + def enable_save_next_step_handlers(self, delay=120, filename="./.handler-saves/step.save"): + """ + Enable saving next step handlers (by default saving disable) + + :param delay: Delay between changes in handlers and saving + :param filename: Filename of save file + """ + self.next_step_saver = Saver(self.next_step_handlers, filename, delay) + + def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"): + """ + Enable saving reply handlers (by default saving disable) + + :param delay: Delay between changes in handlers and saving + :param filename: Filename of save file + """ + self.reply_saver = Saver(self.reply_handlers, filename, delay) + + def disable_save_next_step_handlers(self): + """ + Disable saving next step handlers (by default saving disable) + """ + self.next_step_saver = None + + def disable_save_reply_handlers(self): + """ + Disable saving next step handlers (by default saving disable) + """ + self.reply_saver = None + + def load_next_step_handlers(self, filename="./.handler-saves/step.save", del_file_after_loading=True): + """ + Load next step handlers from save file + + :param filename: Filename of the file where handlers was saved + :param del_file_after_loading: Is passed True, after loading save file will be deleted + """ + self.next_step_saver.load_handlers(filename, del_file_after_loading) + + def load_reply_handlers(self, filename="./.handler-saves/reply.save", del_file_after_loading=True): + """ + Load reply handlers from save file + + :param filename: Filename of the file where handlers was saved + :param del_file_after_loading: Is passed True, after loading save file will be deleted + """ + self.reply_saver.load_handlers(filename) + def set_webhook(self, url=None, certificate=None, max_connections=None, allowed_updates=None): return apihelper.set_webhook(self.token, url, certificate, max_connections, allowed_updates) @@ -984,7 +1102,6 @@ class TeleBot: def get_sticker_set(self, name): """ Use this method to get a sticker set. On success, a StickerSet object is returned. - :param token: :param name: :return: """ @@ -1052,8 +1169,7 @@ class TeleBot: """ Registers a callback function to be notified when a reply to `message` arrives. - Warning: `message` must be sent with reply_markup=types.ForceReply(), otherwise TeleBot will not be able to see - the difference between a reply to `message` and an ordinary message. + Warning: In case `callback` as lambda function, saving reply handlers will not work. :param message: The message for which we are awaiting a reply. :param callback: The callback function to be called when a reply arrives. Must accept one `message` @@ -1066,17 +1182,18 @@ class TeleBot: """ Registers a callback function to be notified when a reply to `message` arrives. - Warning: `message` must be sent with reply_markup=types.ForceReply(), otherwise TeleBot will not be able to see - the difference between a reply to `message` and an ordinary message. + Warning: In case `callback` as lambda function, saving reply handlers will not work. - :param message: The message for which we are awaiting a reply. + :param message_id: The id of the message for which we are awaiting a reply. :param callback: The callback function to be called when a reply arrives. Must accept one `message` parameter, which will contain the replied message. """ if message_id in self.reply_handlers.keys(): - self.reply_handlers[message_id].append({"callback": callback, "args": args, "kwargs": kwargs}) + self.reply_handlers[message_id].append(Handler(callback, *args, **kwargs)) else: - self.reply_handlers[message_id] = [{"callback": callback, "args": args, "kwargs": kwargs}] + self.reply_handlers[message_id] = [Handler(callback, *args, **kwargs)] + if self.reply_saver is not None: + self.reply_saver.start_save_timer() def _notify_reply_handlers(self, new_messages): for message in new_messages: @@ -1087,11 +1204,15 @@ class TeleBot: for handler in handlers: self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"]) self.reply_handlers.pop(reply_msg_id) + if self.reply_saver is not None: + self.reply_saver.start_save_timer() def register_next_step_handler(self, message, callback, *args, **kwargs): """ Registers a callback function to be notified when new message arrives after `message`. + Warning: In case `callback` as lambda function, saving next step handlers will not work. + :param message: The message for which we want to handle new message in the same chat. :param callback: The callback function which next new message arrives. :param args: Args to pass in callback func @@ -1104,15 +1225,20 @@ class TeleBot: """ Registers a callback function to be notified when new message arrives after `message`. + Warning: In case `callback` as lambda function, saving next step handlers will not work. + :param chat_id: The chat for which we want to handle new message. :param callback: The callback function which next new message arrives. :param args: Args to pass in callback func :param kwargs: Args to pass in callback func """ if chat_id in self.next_step_handlers.keys(): - self.next_step_handlers[chat_id].append({"callback": callback, "args": args, "kwargs": kwargs}) + self.next_step_handlers[chat_id].append(Handler(callback, *args, **kwargs)) else: - self.next_step_handlers[chat_id] = [{"callback": callback, "args": args, "kwargs": kwargs}] + self.next_step_handlers[chat_id] = [Handler(callback, *args, **kwargs)] + + if self.next_step_saver is not None: + self.next_step_saver.start_save_timer() def clear_step_handler(self, message): """ @@ -1131,11 +1257,14 @@ class TeleBot: """ self.next_step_handlers[chat_id] = [] + if self.next_step_saver is not None: + self.next_step_saver.start_save_timer() + def clear_reply_handlers(self, message): """ Clears all callback functions registered by register_for_reply() and register_for_reply_by_message_id(). - :param message_id: The message for which we want to clear reply handlers + :param message: The message for which we want to clear reply handlers """ message_id = message.message_id self.clear_reply_handlers_by_message_id(message_id) @@ -1148,6 +1277,9 @@ class TeleBot: """ self.reply_handlers[message_id] = [] + if self.reply_saver is not None: + self.reply_saver.start_save_timer() + def _notify_next_handlers(self, new_messages): i = 0 while i < len(new_messages): @@ -1162,9 +1294,15 @@ class TeleBot: new_messages.pop(i) # removing message that detects with next_step_handler was_poped = True self.next_step_handlers.pop(chat_id, None) + new_messages.pop(i) # removing message that detects with next_step_handler + if self.next_step_saver is not None: + self.next_step_saver.start_save_timer() + i += 1 + if (not was_poped): i += 1 + @staticmethod def _build_handler_dict(handler, **filters): return { @@ -1332,7 +1470,8 @@ class TeleBot: return True - def _test_filter(self, filter, filter_value, message): + @staticmethod + def _test_filter(filter, filter_value, message): test_cases = { 'content_types': lambda msg: msg.content_type in filter_value, 'regexp': lambda msg: msg.content_type == 'text' and re.search(filter_value, msg.text, re.IGNORECASE), @@ -1354,6 +1493,32 @@ class AsyncTeleBot(TeleBot): def __init__(self, *args, **kwargs): TeleBot.__init__(self, *args, **kwargs) + @util.async_dec() + def enable_save_next_step_handlers(self, delay=120, filename="./.handler-saves/step.save", del_file_after_loading=True): + return TeleBot.enable_save_next_step_handlers(self, delay, filename, del_file_after_loading) + + @util.async_dec() + def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save", del_file_after_loading=True): + return TeleBot.enable_save_reply_handlers(self, delay, filename, del_file_after_loading) + + @util.async_dec() + def disable_save_next_step_handlers(self): + return TeleBot.disable_save_next_step_handlers(self) + + @util.async_dec() + def disable_save_reply_handlers(self): + return TeleBot.enable_save_reply_handlers(self) + + @util.async_dec() + def load_next_step_handlers(self, filename="./.handler-saves/step.save"): + return TeleBot.load_next_step_handlers(self, filename) + + @util.async_dec() + def load_reply_handlers(self, filename="./.handler-saves/reply.save"): + return TeleBot.load_reply_handlers(self, filename) + + @util.async_dec() + @util.async_dec() def get_me(self): return TeleBot.get_me(self)