Optimization

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, StateMachine, State
+from telebot.handler_backends import MemoryHandlerBackend, FileHandlerBackend, State
 
 
 REPLY_MARKUP_TYPES = Union[
@@ -188,7 +188,7 @@ class TeleBot:
         self.custom_filters = {}
         self.state_handlers = []
 
-        self.current_states = StateMachine()
+        self.current_states = State()
 
         if apihelper.ENABLE_MIDDLEWARE:
             self.typed_middleware_handlers = {
@@ -2474,7 +2474,7 @@ class TeleBot:
                             if not self._test_filter(message_filter, filter_value, message):
                                 return False
                         need_pop = True
-                        state = State(self.current_states)
+                        state = self.current_states
                         self._exec_task(handler["function"], message, state)
             if need_pop:
                 new_messages.pop(i)  # removing message that was detected by states

telebot/handler_backends.py

@@ -144,7 +144,7 @@ class RedisHandlerBackend(HandlerBackend):
 
 
 
-class StateMachine:
+class State:
     def __init__(self):
         self._states = {}
 
@@ -172,28 +172,20 @@ class StateMachine:
     def _get_data(self, chat_id):
         return self._states[chat_id]['data']
 
-
-class State:
-    """
-    Base class for state managing.
-    """
-    def __init__(self, obj: StateMachine) -> None:
-        self.obj = obj
-
     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.obj.add_state(chat_id,new_state)
+        self.add_state(chat_id,new_state)
 
     def finish(self, chat_id):
         """
         Finish(delete) state of a user.
         :param chat_id:
         """
-        self.obj.delete_state(chat_id)
+        self.delete_state(chat_id)
 
     def retrieve_data(self, chat_id):
         """
@@ -206,13 +198,13 @@ class State:
         Also, at the end of your 'Form' you can get the name:
         data['name']
         """
-        return StateContext(self.obj, chat_id)
+        return StateContext(self, chat_id)
 
 class StateContext:
     """
     Class for data.
     """
-    def __init__(self , obj: StateMachine, chat_id) -> None:
+    def __init__(self , obj: State, chat_id) -> None:
         self.obj = obj
         self.chat_id = chat_id
         self.data = obj._get_data(chat_id)