Add Step/Reply Handler Backend Mechanism

Implement Memory, File, Redis Backends
This commit is contained in:
bedilbek 2020-04-12 22:41:32 +05:00
parent dd726b0759
commit 286188f380
4 changed files with 217 additions and 120 deletions

View File

@ -3,3 +3,4 @@ pytest==3.0.2
requests==2.20.0 requests==2.20.0
six==1.9.0 six==1.9.0
wheel==0.24.0 wheel==0.24.0
redis==3.4.1

View File

@ -20,6 +20,7 @@ setup(name='pyTelegramBotAPI',
install_requires=['requests', 'six'], install_requires=['requests', 'six'],
extras_require={ extras_require={
'json': 'ujson', 'json': 'ujson',
'redis': 'redis>=3.4.1'
}, },
classifiers=[ classifiers=[
'Development Status :: 5 - Production/Stable', 'Development Status :: 5 - Production/Stable',

View File

@ -2,8 +2,6 @@
from __future__ import print_function from __future__ import print_function
import logging import logging
import os
import pickle
import re import re
import sys import sys
import threading import threading
@ -23,6 +21,7 @@ logger.addHandler(console_output_handler)
logger.setLevel(logging.ERROR) logger.setLevel(logging.ERROR)
from telebot import apihelper, types, util from telebot import apihelper, types, util
from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend
""" """
Module : telebot Module : telebot
@ -43,64 +42,6 @@ class Handler:
return getattr(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:
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: class TeleBot:
""" This is TeleBot Class """ This is TeleBot Class
Methods: Methods:
@ -141,7 +82,10 @@ class TeleBot:
answerInlineQuery 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 :param token: bot API token
:return: Telebot object. :return: Telebot object.
@ -155,14 +99,13 @@ class TeleBot:
self.last_update_id = 0 self.last_update_id = 0
self.exc_info = None self.exc_info = None
# key: message_id, value: handler list self.next_step_backend = next_step_backend
self.reply_handlers = {} if not self.next_step_backend:
self.next_step_backend = MemoryHandlerBackend()
# key: chat_id, value: handler list self.reply_backend = reply_backend
self.next_step_handlers = {} if not self.reply_backend:
self.reply_backend = MemoryHandlerBackend()
self.next_step_saver = None
self.reply_saver = None
self.message_handlers = [] self.message_handlers = []
self.edited_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"): 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 delay: Delay between changes in handlers and saving
:param filename: Filename of save file :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"): def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"):
""" """
Enable saving reply handlers (by default saving disable) 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 delay: Delay between changes in handlers and saving
:param filename: Filename of save file :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): def disable_save_next_step_handlers(self):
""" """
Disable saving next step handlers (by default saving disable) 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): def disable_save_reply_handlers(self):
""" """
Disable saving next step handlers (by default saving disable) 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): def load_next_step_handlers(self, filename="./.handler-saves/step.save", del_file_after_loading=True):
""" """
Load next step handlers from save file 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 filename: Filename of the file where handlers was saved
:param del_file_after_loading: Is passed True, after loading save file will be deleted :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): def load_reply_handlers(self, filename="./.handler-saves/reply.save", del_file_after_loading=True):
""" """
Load reply handlers from save file 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 filename: Filename of the file where handlers was saved
:param del_file_after_loading: Is passed True, after loading save file will be deleted :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): 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) 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` :param callback: The callback function to be called when a reply arrives. Must accept one `message`
parameter, which will contain the replied message. parameter, which will contain the replied message.
""" """
if message_id in self.reply_handlers.keys(): self.reply_backend.register_handler(message_id, Handler(callback, *args, **kwargs))
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()
def _notify_reply_handlers(self, new_messages): def _notify_reply_handlers(self, new_messages):
""" """
@ -1414,14 +1390,9 @@ class TeleBot:
""" """
for message in new_messages: for message in new_messages:
if hasattr(message, "reply_to_message") and message.reply_to_message is not None: if hasattr(message, "reply_to_message") and message.reply_to_message is not None:
reply_msg_id = message.reply_to_message.message_id handlers = self.reply_backend.get_handlers(message.reply_to_message.message_id)
if reply_msg_id in self.reply_handlers.keys(): for handler in handlers:
handlers = self.reply_handlers[reply_msg_id] self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"])
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): 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 args: Args to pass in callback func
:param kwargs: 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_backend.register_handler(chat_id, Handler(callback, *args, **kwargs))
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()
def clear_step_handler(self, message): 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 :param chat_id: The chat for which we want to clear next step handlers
""" """
self.next_step_handlers[chat_id] = [] self.next_step_backend.clear_handlers(chat_id)
if self.next_step_saver is not None:
self.next_step_saver.start_save_timer()
def clear_reply_handlers(self, message): 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 :param message_id: The message id for which we want to clear reply handlers
""" """
self.reply_handlers[message_id] = [] self.reply_backend.clear_handlers(message_id)
if self.reply_saver is not None:
self.reply_saver.start_save_timer()
def _notify_next_handlers(self, new_messages): def _notify_next_handlers(self, new_messages):
""" """
@ -1502,22 +1461,14 @@ class TeleBot:
:param new_messages: :param new_messages:
:return: :return:
""" """
i = 0 for i, message in enumerate(new_messages):
while i < len(new_messages): need_pop = False
message = new_messages[i] handlers = self.next_step_backend.get_handlers(message.chat.id)
chat_id = message.chat.id for handler in handlers:
was_poped = False need_pop = True
if chat_id in self.next_step_handlers.keys(): self._exec_task(handler["callback"], message, *handler["args"], **handler["kwargs"])
handlers = self.next_step_handlers.pop(chat_id, None) if need_pop:
if handlers: new_messages.pop(i) # removing message that detects with next_step_handler
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
@staticmethod @staticmethod
def _build_handler_dict(handler, **filters): def _build_handler_dict(handler, **filters):

144
telebot/handler_backends.py Normal file
View File

@ -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