diff --git a/requirements.txt b/requirements.txt index 6e4ca40..f58156a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ pytest==3.0.2 requests==2.20.0 six==1.9.0 wheel==0.24.0 +redis==3.4.1 diff --git a/setup.py b/setup.py index 44a5cde..7c7c856 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ setup(name='pyTelegramBotAPI', install_requires=['requests', 'six'], extras_require={ 'json': 'ujson', + 'redis': 'redis>=3.4.1' }, classifiers=[ 'Development Status :: 5 - Production/Stable', diff --git a/telebot/__init__.py b/telebot/__init__.py index 72f6c71..89b792a 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -2,8 +2,6 @@ from __future__ import print_function import logging -import os -import pickle import re import sys import threading @@ -23,6 +21,7 @@ logger.addHandler(console_output_handler) logger.setLevel(logging.ERROR) from telebot import apihelper, types, util +from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend """ Module : telebot @@ -43,64 +42,6 @@ class Handler: 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: - if (apihelper.CUSTOM_SERIALIZER is None): - pickle.dump(handlers, file) - else: - apihelper.CUSTOM_SERIALIZER.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: - if (apihelper.CUSTOM_SERIALIZER is None): - handlers = pickle.load(file) - else: - handlers = apihelper.CUSTOM_SERIALIZER.load(file) - - if del_file_after_loading: - os.remove(filename) - - return handlers - - class TeleBot: """ This is TeleBot Class Methods: @@ -141,7 +82,10 @@ class TeleBot: answerInlineQuery """ - def __init__(self, token, threaded=True, skip_pending=False, num_threads=2): + def __init__( + self, token, threaded=True, skip_pending=False, num_threads=2, + next_step_backend=None, reply_backend=None + ): """ :param token: bot API token :return: Telebot object. @@ -155,14 +99,13 @@ class TeleBot: self.last_update_id = 0 self.exc_info = None - # key: message_id, value: handler list - self.reply_handlers = {} + self.next_step_backend = next_step_backend + if not self.next_step_backend: + self.next_step_backend = MemoryHandlerBackend() - # key: chat_id, value: handler list - self.next_step_handlers = {} - - self.next_step_saver = None - self.reply_saver = None + self.reply_backend = reply_backend + if not self.reply_backend: + self.reply_backend = MemoryHandlerBackend() self.message_handlers = [] self.edited_message_handlers = [] @@ -196,51 +139,89 @@ class TeleBot: def enable_save_next_step_handlers(self, delay=120, filename="./.handler-saves/step.save"): """ - Enable saving next step handlers (by default saving disable) + Enable saving next step handlers (by default saving disabled) + + This function explicitly assigns FileHandlerBackend (instead of Saver) just to keep backward + compatibility whose purpose was to enable file saving capability for handlers. And the same + implementation is now available with FileHandlerBackend + + Most probably this function should be deprecated in future major releases :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) + self.next_step_backend = FileHandlerBackend(self.next_step_backend.handlers, filename, delay) def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"): """ Enable saving reply handlers (by default saving disable) + This function explicitly assigns FileHandlerBackend (instead of Saver) just to keep backward + compatibility whose purpose was to enable file saving capability for handlers. And the same + implementation is now available with FileHandlerBackend + + Most probably this function should be deprecated in future major releases + :param delay: Delay between changes in handlers and saving :param filename: Filename of save file """ - self.reply_saver = Saver(self.reply_handlers, filename, delay) + self.reply_backend = FileHandlerBackend(self.reply_backend.handlers, filename, delay) def disable_save_next_step_handlers(self): """ Disable saving next step handlers (by default saving disable) + + This function is left to keep backward compatibility whose purpose was to disable file saving capability + for handlers. For the same purpose, MemoryHandlerBackend is reassigned as a new next_step_backend backend + instead of FileHandlerBackend. + + Most probably this function should be deprecated in future major releases """ - self.next_step_saver = None + self.next_step_backend = MemoryHandlerBackend(self.next_step_backend.handlers) def disable_save_reply_handlers(self): """ Disable saving next step handlers (by default saving disable) + + This function is left to keep backward compatibility whose purpose was to disable file saving capability + for handlers. For the same purpose, MemoryHandlerBackend is reassigned as a new reply_backend backend + instead of FileHandlerBackend. + + Most probably this function should be deprecated in future major releases """ - self.reply_saver = None + self.reply_backend = MemoryHandlerBackend(self.reply_backend.handlers) def load_next_step_handlers(self, filename="./.handler-saves/step.save", del_file_after_loading=True): """ Load next step handlers from save file + This function is left to keep backward compatibility whose purpose was to load handlers from file with the + help of FileHandlerBackend and is only recommended to use if next_step_backend was assigned as + FileHandlerBackend before entering this function + + Most probably this function should be deprecated in future major releases + :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) + self.next_step_backend: FileHandlerBackend + self.next_step_backend.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 + This function is left to keep backward compatibility whose purpose was to load handlers from file with the + help of FileHandlerBackend and is only recommended to use if reply_backend was assigned as + FileHandlerBackend before entering this function + + Most probably this function should be deprecated in future major releases + :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) + self.reply_backend: FileHandlerBackend + self.reply_backend.load_handlers(filename, del_file_after_loading) 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) @@ -1399,12 +1380,7 @@ class TeleBot: :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(Handler(callback, *args, **kwargs)) - else: - self.reply_handlers[message_id] = [Handler(callback, *args, **kwargs)] - if self.reply_saver is not None: - self.reply_saver.start_save_timer() + self.reply_backend.register_handler(message_id, Handler(callback, *args, **kwargs)) def _notify_reply_handlers(self, new_messages): """ @@ -1414,14 +1390,9 @@ class TeleBot: """ for message in new_messages: if hasattr(message, "reply_to_message") and message.reply_to_message is not None: - reply_msg_id = message.reply_to_message.message_id - if reply_msg_id in self.reply_handlers.keys(): - handlers = self.reply_handlers[reply_msg_id] - 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() + handlers = self.reply_backend.get_handlers(message.reply_to_message.message_id) + for handler in handlers: + self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"]) def register_next_step_handler(self, message, callback, *args, **kwargs): """ @@ -1448,13 +1419,7 @@ class TeleBot: :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(Handler(callback, *args, **kwargs)) - else: - 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() + self.next_step_backend.register_handler(chat_id, Handler(callback, *args, **kwargs)) def clear_step_handler(self, message): """ @@ -1471,10 +1436,7 @@ class TeleBot: :param chat_id: The chat for which we want to clear next step handlers """ - self.next_step_handlers[chat_id] = [] - - if self.next_step_saver is not None: - self.next_step_saver.start_save_timer() + self.next_step_backend.clear_handlers(chat_id) def clear_reply_handlers(self, message): """ @@ -1491,10 +1453,7 @@ class TeleBot: :param message_id: The message id for which we want to clear reply handlers """ - self.reply_handlers[message_id] = [] - - if self.reply_saver is not None: - self.reply_saver.start_save_timer() + self.reply_backend.clear_handlers(message_id) def _notify_next_handlers(self, new_messages): """ @@ -1502,22 +1461,14 @@ class TeleBot: :param new_messages: :return: """ - i = 0 - while i < len(new_messages): - message = new_messages[i] - chat_id = message.chat.id - was_poped = False - if chat_id in self.next_step_handlers.keys(): - handlers = self.next_step_handlers.pop(chat_id, None) - if handlers: - for handler in handlers: - self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"]) - new_messages.pop(i) # removing message that detects with next_step_handler - was_poped = True - if self.next_step_saver is not None: - self.next_step_saver.start_save_timer() - if not was_poped: - i += 1 + for i, message in enumerate(new_messages): + need_pop = False + handlers = self.next_step_backend.get_handlers(message.chat.id) + for handler in handlers: + need_pop = True + self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"]) + if need_pop: + new_messages.pop(i) # removing message that detects with next_step_handler @staticmethod def _build_handler_dict(handler, **filters): diff --git a/telebot/handler_backends.py b/telebot/handler_backends.py new file mode 100644 index 0000000..97dd44f --- /dev/null +++ b/telebot/handler_backends.py @@ -0,0 +1,144 @@ +import os +import pickle +import threading + +from telebot import apihelper + + +class HandlerBackend: + """ + Class for saving (next step|reply) handlers + """ + handlers = {} + + def __init__(self, handlers=None): + if handlers: + self.handlers = handlers + + def register_handler(self, handler_group_id, handler): + raise NotImplementedError() + + def clear_handlers(self, handler_group_id): + raise NotImplementedError() + + def get_handlers(self, handler_group_id): + raise NotImplementedError() + + +class MemoryHandlerBackend(HandlerBackend): + def register_handler(self, handler_group_id, handler): + if handler_group_id in self.handlers: + self.handlers[handler_group_id].append(handler) + else: + self.handlers[handler_group_id] = [handler] + + def clear_handlers(self, handler_group_id): + self.handlers.pop(handler_group_id, []) + + def get_handlers(self, handler_group_id): + return self.handlers.pop(handler_group_id, []) + + +class FileHandlerBackend(HandlerBackend): + def __init__(self, handlers=None, filename='./.handler-saves/handlers.save', delay=120): + super().__init__(handlers) + self.filename = filename + self.delay = delay + self.timer = threading.Timer(delay, self.save_handlers) + + def register_handler(self, handler_group_id, handler): + if handler_group_id in self.handlers: + self.handlers[handler_group_id].append(handler) + else: + self.handlers[handler_group_id] = [handler] + + self.start_save_timer() + + def clear_handlers(self, handler_group_id): + self.handlers.pop(handler_group_id, []) + + self.start_save_timer() + + def get_handlers(self, handler_group_id): + handlers = self.handlers.pop(handler_group_id, []) + + self.start_save_timer() + + return 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: + if (apihelper.CUSTOM_SERIALIZER is None): + pickle.dump(handlers, file) + else: + apihelper.CUSTOM_SERIALIZER.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: + if (apihelper.CUSTOM_SERIALIZER is None): + handlers = pickle.load(file) + else: + handlers = apihelper.CUSTOM_SERIALIZER.load(file) + + if del_file_after_loading: + os.remove(filename) + + return handlers + + +class RedisHandlerBackend(HandlerBackend): + def __init__(self, handlers=None, host='localhost', port=6379, db=0, prefix='telebot'): + super().__init__(handlers) + from redis import Redis + self.prefix = prefix + self.redis = Redis(host, port, db) + + def _key(self, handle_group_id): + return ':'.join((self.prefix, str(handle_group_id))) + + def register_handler(self, handler_group_id, handler): + handlers = [] + value = self.redis.get(self._key(handler_group_id)) + if value: + handlers = pickle.loads(value) + handlers.append(handler) + self.redis.set(self._key(handler_group_id), pickle.dumps(handlers)) + + def clear_handlers(self, handler_group_id): + self.redis.delete(self._key(handler_group_id)) + + def get_handlers(self, handler_group_id): + handlers = [] + value = self.redis.get(self._key(handler_group_id)) + if value: + handlers = pickle.loads(value) + self.clear_handlers(handler_group_id) + + return handlers