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
six==1.9.0
wheel==0.24.0
redis==3.4.1

View File

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

View File

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

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