Merge branch 'master' of https://github.com/eternnoir/pyTelegramBotAPI

telebot/__init__.py

@@ -1,13 +1,15 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function
 
-import threading
-import time
+import logging
+import os
+import pickle
 import re
 import sys
-import six
+import threading
+import time
 
-import logging
+import six
 
 logger = logging.getLogger('TeleBot')
 formatter = logging.Formatter(
@@ -27,6 +29,72 @@ 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 +153,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 +170,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)
 
@@ -880,6 +999,12 @@ class TeleBot:
             return result
         return types.Message.de_json(result)
 
+    def edit_message_media(self, media, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
+        result = apihelper.edit_message_media(self.token, media, chat_id, message_id, inline_message_id, reply_markup)
+        if type(result) == bool:  # if edit inline message return is bool not Message.
+            return result
+        return types.Message.de_json(result)
+
     def edit_message_reply_markup(self, chat_id=None, message_id=None, inline_message_id=None, reply_markup=None):
         result = apihelper.edit_message_reply_markup(self.token, chat_id, message_id, inline_message_id, reply_markup)
         if type(result) == bool:
@@ -984,7 +1109,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 +1176,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 +1189,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 +1211,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 +1232,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 +1264,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 +1284,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):
@@ -1155,21 +1294,22 @@ class TeleBot:
             chat_id = message.chat.id
             was_poped = False
             if chat_id in self.next_step_handlers.keys():
-                handlers = self.next_step_handlers[chat_id]
-                if (handlers):
+                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
-                self.next_step_handlers.pop(chat_id, None)
-            if (not was_poped):
+                if self.next_step_saver is not None:
+                    self.next_step_saver.start_save_timer()
+            if not was_poped:
                 i += 1
 
     @staticmethod
     def _build_handler_dict(handler, **filters):
         return {
             'function': handler,
-            'filters': filters
+            'filters' : filters
         }
 
     def message_handler(self, commands=None, regexp=None, func=None, content_types=['text'], **kwargs):
@@ -1332,7 +1472,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 +1495,30 @@ 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"):
+        return TeleBot.enable_save_next_step_handlers(self, delay, filename)
+
+    @util.async_dec()
+    def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"):
+        return TeleBot.enable_save_reply_handlers(self, delay, filename)
+
+    @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", del_file_after_loading=True):
+        return TeleBot.load_next_step_handlers(self, filename, del_file_after_loading)
+
+    @util.async_dec()
+    def load_reply_handlers(self, filename="./.handler-saves/reply.save", del_file_after_loading=True):
+        return TeleBot.load_reply_handlers(self, filename, del_file_after_loading)
+
     @util.async_dec()
     def get_me(self):
         return TeleBot.get_me(self)
@@ -1514,6 +1679,10 @@ class AsyncTeleBot(TeleBot):
     def edit_message_text(self, *args, **kwargs):
         return TeleBot.edit_message_text(self, *args, **kwargs)
 
+    @util.async_dec()
+    def edit_message_media(self, *args, **kwargs):
+        return TeleBot.edit_message_media(self, *args, **kwargs)
+
     @util.async_dec()
     def edit_message_reply_markup(self, *args, **kwargs):
         return TeleBot.edit_message_reply_markup(self, *args, **kwargs)