Merge pull request #1336 from coder2020official/master

File support for states

examples/custom_states.py

@@ -61,4 +61,8 @@ def age_incorrect(message):
 
 bot.add_custom_filter(custom_filters.StateFilter(bot))
 bot.add_custom_filter(custom_filters.IsDigitFilter())
+
+# set saving states into file.
+bot.enable_saving_states() # you can delete this if you do not need to save states
+
 bot.infinity_polling(skip_pending=True)

telebot/__init__.py

@@ -27,7 +27,7 @@ logger.addHandler(console_output_handler)
 logger.setLevel(logging.ERROR)
 
 from telebot import apihelper, util, types
-from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend, State
+from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend, StateMemory, StateFile
 
 
 REPLY_MARKUP_TYPES = Union[
@@ -188,7 +188,8 @@ class TeleBot:
         self.custom_filters = {}
         self.state_handlers = []
 
-        self.current_states = State()
+        self.current_states = StateMemory()
+
 
         if apihelper.ENABLE_MIDDLEWARE:
             self.typed_middleware_handlers = {
@@ -237,6 +238,17 @@ class TeleBot:
         """
         self.next_step_backend = FileHandlerBackend(self.next_step_backend.handlers, filename, delay)
 
+    def enable_saving_states(self, filename="./.state-save/states.pkl"):
+        """
+        Enable saving states (by default saving disabled)
+
+        :param filename: Filename of saving file
+
+        """
+
+        self.current_states = StateFile(filename=filename)
+        self.current_states._create_dir()
+
     def enable_save_reply_handlers(self, delay=120, filename="./.handler-saves/reply.save"):
         """
         Enable saving reply handlers (by default saving disable)
@@ -345,7 +357,7 @@ class TeleBot:
         """
         return apihelper.delete_webhook(self.token, drop_pending_updates, timeout)
 
-    def get_webhook_info(self, timeout=None):
+    def get_webhook_info(self, timeout: Optional[int]=None):
         """
         Use this method to get current webhook status. Requires no parameters.
         If the bot is using getUpdates, will return an object with the url field empty.
@@ -2369,7 +2381,7 @@ class TeleBot:
         chat_id = message.chat.id
         self.register_next_step_handler_by_chat_id(chat_id, callback, *args, **kwargs)
 
-    def set_state(self, chat_id, state):
+    def set_state(self, chat_id: int, state: Union[int, str]):
         """
         Sets a new state of a user.
         :param chat_id:
@@ -2377,7 +2389,7 @@ class TeleBot:
         """
         self.current_states.add_state(chat_id, state)
 
-    def delete_state(self, chat_id):
+    def delete_state(self, chat_id: int):
         """
         Delete the current state of a user.
         :param chat_id:
@@ -2385,10 +2397,10 @@ class TeleBot:
         """
         self.current_states.delete_state(chat_id)
 
-    def retrieve_data(self, chat_id):
+    def retrieve_data(self, chat_id: int):
         return self.current_states.retrieve_data(chat_id)
 
-    def get_state(self, chat_id):
+    def get_state(self, chat_id: int):
         """
         Get current state of a user.
         :param chat_id:
@@ -2396,7 +2408,7 @@ class TeleBot:
         """
         return self.current_states.current_state(chat_id)
 
-    def add_data(self, chat_id, **kwargs):
+    def add_data(self, chat_id: int, **kwargs):
         """
         Add data to states.
         :param chat_id:

telebot/custom_filters.py

@@ -158,9 +158,9 @@ class StateFilter(AdvancedCustomFilter):
     key = 'state'
 
     def check(self, message, text):
-        if self.bot.current_states.current_state(message.from_user.id) is False:return False
-        elif text == '*':return True
-        elif type(text) is list:return self.bot.current_states.current_state(message.from_user.id) in text
+        if self.bot.current_states.current_state(message.from_user.id) is False: return False
+        elif text == '*': return True
+        elif type(text) is list: return self.bot.current_states.current_state(message.from_user.id) in text
         return self.bot.current_states.current_state(message.from_user.id) == text
 
 class IsDigitFilter(SimpleCustomFilter):

telebot/handler_backends.py

@@ -143,7 +143,7 @@ class RedisHandlerBackend(HandlerBackend):
         return handlers
 
 
-class State:
+class StateMemory:
     def __init__(self):
         self._states = {}
 
@@ -166,7 +166,7 @@ class State:
 
     def delete_state(self, chat_id):
         """Delete a state"""
-        return self._states.pop(chat_id)
+        self._states.pop(chat_id)
 
     def _get_data(self, chat_id):
         return self._states[chat_id]['data']
@@ -195,7 +195,7 @@ class State:
         Save input text.
 
         Usage:
-        with state.retrieve_data(message.chat.id) as data:
+        with bot.retrieve_data(message.chat.id) as data:
             data['name'] = message.text
 
         Also, at the end of your 'Form' you can get the name:
@@ -203,11 +203,114 @@ class State:
         """
         return StateContext(self, chat_id)
 
+
+class StateFile:
+    """
+    Class to save states in a file.
+    """
+    def __init__(self, filename):
+        self.file_path = filename
+
+    def add_state(self, chat_id, state):
+        """
+        Add a state.
+        :param chat_id:
+        :param state: new state
+        """
+        states_data = self._read_data()
+        if chat_id in states_data:
+            states_data[chat_id]['state'] = state
+            return self._save_data(states_data)
+        else:
+            new_data = states_data[chat_id] = {'state': state,'data': {}}
+            return self._save_data(states_data)
+
+
+    def current_state(self, chat_id):
+        """Current state."""
+        states_data = self._read_data()
+        if chat_id in states_data: return states_data[chat_id]['state']
+        else: return False
+
+    def delete_state(self, chat_id):
+        """Delete a state"""
+        states_data = self._read_data()
+        states_data.pop(chat_id)
+        self._save_data(states_data)
+
+    def _read_data(self):
+        """
+        Read the data from file.
+        """
+        file = open(self.file_path, 'rb')
+        states_data = pickle.load(file)
+        file.close()
+        return states_data
+
+    def _create_dir(self):
+        """
+        Create directory .save-handlers.
+        """
+        dirs = self.file_path.rsplit('/', maxsplit=1)[0]
+        os.makedirs(dirs, exist_ok=True)
+        if not os.path.isfile(self.file_path):
+            with open(self.file_path,'wb') as file:
+                pickle.dump({}, file)
+        
+    def _save_data(self, new_data):
+        """
+        Save data after editing.
+        :param new_data:
+        """
+        with open(self.file_path, 'wb+') as state_file:
+            pickle.dump(new_data, state_file, protocol=pickle.HIGHEST_PROTOCOL)
+        return True
+
+    def _get_data(self, chat_id):
+        return self._read_data()[chat_id]['data']
+
+    def set(self, chat_id, new_state):
+        """
+        Set a new state for a user.
+        :param chat_id:
+        :param new_state: new_state of a user
+        
+        """
+        self.add_state(chat_id,new_state)
+
+    def _add_data(self, chat_id, key, value):
+        states_data = self._read_data()
+        result = states_data[chat_id]['data'][key] = value
+        self._save_data(result)
+
+        return result
+
+    def finish(self, chat_id):
+        """
+        Finish(delete) state of a user.
+        :param chat_id:
+        """
+        self.delete_state(chat_id)
+
+    def retrieve_data(self, chat_id):
+        """
+        Save input text.
+
+        Usage:
+        with bot.retrieve_data(message.chat.id) as data:
+            data['name'] = message.text
+
+        Also, at the end of your 'Form' you can get the name:
+        data['name']
+        """
+        return StateFileContext(self, chat_id)
+
+
 class StateContext:
     """
     Class for data.
     """
-    def __init__(self , obj: State, chat_id) -> None:
+    def __init__(self , obj: StateMemory, chat_id) -> None:
         self.obj = obj
         self.chat_id = chat_id
         self.data = obj._get_data(chat_id)
@@ -217,3 +320,23 @@ class StateContext:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return
+
+class StateFileContext:
+    """
+    Class for data.
+    """
+    def __init__(self , obj: StateFile, chat_id) -> None:
+        self.obj = obj
+        self.chat_id = chat_id
+        self.data = self.obj._get_data(self.chat_id)
+
+    def __enter__(self):
+        return self.data
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        old_data = self.obj._read_data()
+        for i in self.data:
+            old_data[self.chat_id]['data'][i] = self.data.get(i)
+        self.obj._save_data(old_data)
+
+        return