Sync states v2 early version

Author: coder2020official

telebot/handler_backends.py

@@ -9,6 +9,8 @@
 except:
     redis_installed = False
 
+# backward compatibility
+from telebot.states import State, StatesGroup
 
 class HandlerBackend(object):
     """
@@ -160,45 +162,6 @@ def get_handlers(self, handler_group_id):
         return handlers
 
 
-class State:
-    """
-    Class representing a state.
-
-    .. code-block:: python3
-
-        class MyStates(StatesGroup):
-            my_state = State() # returns my_state:State string.
-    """
-    def __init__(self) -> None:
-        self.name = None
-    def __str__(self) -> str:
-        return self.name
-
-
-class StatesGroup:
-    """
-    Class representing common states.
-
-    .. code-block:: python3
-
-        class MyStates(StatesGroup):
-            my_state = State() # returns my_state:State string.
-    """
-    def __init_subclass__(cls) -> None:
-        state_list = []
-        for name, value in cls.__dict__.items():
-            if not name.startswith('__') and not callable(value) and isinstance(value, State):
-                # change value of that variable
-                value.name = ':'.join((cls.__name__, name))
-                value.group = cls
-                state_list.append(value)
-        cls._state_list = state_list
-
-    @classmethod
-    def state_list(self):
-        return self._state_list
-
-
 class BaseMiddleware:
     """
     Base class for middleware.
@@ -292,4 +255,4 @@ def start2(message):
     
     """
     def __init__(self) -> None:
-        pass
+        pass

telebot/states/__init__.py

@@ -0,0 +1,43 @@
+"""
+Contains classes for states and state groups.
+"""
+
+
+class State:
+    """
+    Class representing a state.
+
+    .. code-block:: python3
+
+        class MyStates(StatesGroup):
+            my_state = State() # returns my_state:State string.
+    """
+    def __init__(self) -> None:
+        self.name: str = None
+        self.group: StatesGroup = None
+    def __str__(self) -> str:
+        return f"<{self.group.__name__}:{self.name}>"
+
+
+class StatesGroup:
+    """
+    Class representing common states.
+
+    .. code-block:: python3
+
+        class MyStates(StatesGroup):
+            my_state = State() # returns my_state:State string.
+    """
+    def __init_subclass__(cls) -> None:
+        state_list = []
+        for name, value in cls.__dict__.items():
+            if not name.startswith('__') and not callable(value) and isinstance(value, State):
+                # change value of that variable
+                value.name = ':'.join((cls.__name__, name))
+                value.group = cls
+                state_list.append(value)
+        cls._state_list = state_list
+
+    @classmethod
+    def state_list(self):
+        return self._state_list

telebot/states/aio/__init__.py

@@ -0,0 +1,7 @@
+from .context import StateContext
+from .middleware import StateMiddleware
+
+__all__ = [
+    'StateContext',
+    'StateMiddleware',
+]

telebot/states/sync/__init__.py

@@ -0,0 +1,153 @@
+from telebot.states import State, StatesGroup
+from telebot.types import CallbackQuery, Message
+from telebot import TeleBot
+
+from typing import Union
+
+
+class StateContext():
+    """
+    Class representing a state context.
+
+    Passed through a middleware to provide easy way to set states.
+
+    .. code-block:: python3
+
+        @bot.message_handler(commands=['start'])
+        def start_ex(message: types.Message, state_context: StateContext):
+            state_context.set(MyStates.name)
+            bot.send_message(message.chat.id, 'Hi, write me a name', reply_to_message_id=message.message_id)
+            # also, state_context.data(), .add_data(), .reset_data(), .delete() methods available.
+    """
+    
+    def __init__(self, message: Union[Message, CallbackQuery], bot: str) -> None:
+        self.message: Union[Message, CallbackQuery] = message
+        self.bot: TeleBot = bot
+        self.bot_id = self.bot.bot_id
+
+    def _resolve_context(self) -> Union[Message, CallbackQuery]:
+        chat_id = None
+        user_id = None
+        business_connection_id = self.message.business_connection_id
+        bot_id = self.bot_id
+        message_thread_id = None
+
+        if isinstance(self.message, Message):
+            chat_id = self.message.chat.id
+            user_id = self.message.from_user.id
+            message_thread_id = self.message.message_thread_id if self.message.is_topic_message else None
+        elif isinstance(self.message, CallbackQuery):
+            chat_id = self.message.message.chat.id
+            user_id = self.message.from_user.id
+            message_thread_id = self.message.message.message_thread_id if self.message.message.is_topic_message else None
+
+        return chat_id, user_id, business_connection_id, bot_id, message_thread_id
+    
+    def set(self, state: Union[State, str]) -> None:
+        """
+        Set state for current user.
+
+        :param state: State object or state name.
+        :type state: Union[State, str]
+
+        .. code-block:: python3
+
+            @bot.message_handler(commands=['start'])
+            def start_ex(message: types.Message, state_context: StateContext):
+                state_context.set(MyStates.name)
+                bot.send_message(message.chat.id, 'Hi, write me a name', reply_to_message_id=message.message_id)
+        """
+
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        if isinstance(state, State):
+            state = state.name
+        return self.bot.set_state(
+            chat_id=chat_id,
+            user_id=user_id,
+            state=state,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id
+        )
+    
+    def get(self) -> str:
+        """
+        Get current state for current user.
+
+        :return: Current state name.
+        :rtype: str
+        """
+
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        return self.bot.get_state(
+            chat_id=chat_id,
+            user_id=user_id,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id
+        )
+    
+    def reset_data(self) -> None:
+        """
+        Reset data for current user. 
+        State will not be changed.
+        """
+
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        return self.bot.reset_data(
+            chat_id=chat_id,
+            user_id=user_id,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id
+        )
+    
+    def delete(self) -> None:
+        """
+        Deletes state and data for current user.
+        """
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        return self.bot.delete_state(
+            chat_id=chat_id,
+            user_id=user_id,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id
+        )
+    
+    def data(self) -> dict:
+        """
+        Get data for current user.
+
+        .. code-block:: python3
+ 
+            with state_context.data() as data:
+                print(data)
+        """
+
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        return self.bot.retrieve_data(
+            chat_id=chat_id,
+            user_id=user_id,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id
+        )
+    
+    def add_data(self, **kwargs) -> None:
+        """
+        Add data for current user.
+
+        :param kwargs: Data to add.
+        :type kwargs: dict
+        """
+
+        chat_id, user_id, business_connection_id, bot_id, message_thread_id = self._resolve_context()
+        return self.bot.add_data(
+            chat_id=chat_id,
+            user_id=user_id,
+            business_connection_id=business_connection_id,
+            bot_id=bot_id,
+            message_thread_id=message_thread_id,
+            **kwargs
+        )

telebot/states/sync/context.py

@@ -0,0 +1,17 @@
+from telebot.handler_backends import BaseMiddleware
+from telebot import TeleBot
+from telebot.states.sync.context import StateContext
+
+
+class StateMiddleware(BaseMiddleware):
+
+    def __init__(self, bot: TeleBot) -> None:
+        self.update_sensitive = False
+        self.update_types = ['message', 'edited_message', 'callback_query'] #TODO: support other types
+        self.bot: TeleBot = bot
+
+    def pre_process(self, message, data):
+        data['state_context'] = StateContext(message, self.bot)
+
+    def post_process(self, message, data, exception):
+        pass