From 64347d856267c271887a783a54de1650237ca63b Mon Sep 17 00:00:00 2001
From: Hiddify <114227601+hiddify-com@users.noreply.github.com>
Date: Fri, 12 Jul 2024 10:23:31 +0000
Subject: [PATCH] implement  step_handler and register_next_step_handler

---
 hiddifypanel_bot/basebot/bot.py            | 14 ++---
 hiddifypanel_bot/basebot/hasync_telebot.py | 67 ++++++++++++++++++++++
 requirements.txt                           |  4 +-
 3 files changed, 73 insertions(+), 12 deletions(-)
 create mode 100644 hiddifypanel_bot/basebot/hasync_telebot.py

diff --git a/hiddifypanel_bot/basebot/bot.py b/hiddifypanel_bot/basebot/bot.py
index b3761d3..c357a98 100644
--- a/hiddifypanel_bot/basebot/bot.py
+++ b/hiddifypanel_bot/basebot/bot.py
@@ -1,6 +1,6 @@
 import logging
 import os
-from telebot.async_telebot import AsyncTeleBot,logger
+from telebot.async_telebot import logger
 from telebot import TeleBot, ExceptionHandler,types
 from telebot import asyncio_filters
 from telebot.asyncio_storage import StateMemoryStorage
@@ -8,6 +8,7 @@
 from .call_action_filter import CallActionFilter
 from .role_filter import RoleFilter
 from .middleware import Middleware
+from .hasync_telebot  import HAsyncTeleBot
 logger.setLevel(logging.INFO)  # Outputs debug messages to console.
 
 
@@ -21,19 +22,12 @@ def handle(self, exception):
         logger.error(exception)
         raise exception
 
-bot: AsyncTeleBot = AsyncTeleBot(
+bot: HAsyncTeleBot = HAsyncTeleBot(
     BOT_TOKEN,
-    # exception_handler=MyExceptionHandler(),
-    
+    exception_handler=MyExceptionHandler(),
     state_storage=state_storage,
 )
 
-def callback(func=None,**kwargs):
-    if func is None: func=lambda call:True
-    return bot.callback_query_handler_orig(func,**kwargs)
-bot.callback_query_handler_orig=bot.callback_query_handler
-bot.callback_query_handler=callback
-
 
 bot.add_custom_filter(asyncio_filters.StateFilter(bot))
 bot.add_custom_filter(asyncio_filters.TextMatchFilter())
diff --git a/hiddifypanel_bot/basebot/hasync_telebot.py b/hiddifypanel_bot/basebot/hasync_telebot.py
new file mode 100644
index 0000000..03131df
--- /dev/null
+++ b/hiddifypanel_bot/basebot/hasync_telebot.py
@@ -0,0 +1,67 @@
+from typing import Callable
+from telebot.async_telebot import AsyncTeleBot,logger
+
+class HAsyncTeleBot(AsyncTeleBot):
+
+    def step_handler(self):
+        """
+            This decorator should be on top of a method for being used as step handeler
+
+            Example: 
+            @bot.step_handler()
+            def test(message):
+                pass
+        """
+        def decorator(handler):
+            handler_name = handler.__name__
+            module_name = handler.__module__
+            signature = f"{hash(handler.__code__.co_filename)}:{module_name}.{handler_name}:{handler.__code__.co_firstlineno}"
+            async def wrapper(*args, **kwargs):
+                msg=args[0]
+                async with self.retrieve_data(msg.from_user.id,msg.chat.id) as data:
+                    step_kwargs=data.get('__step_handler_kwargs__',{})
+                    step_args=data.get('__step_handler_args__',())
+                    if '__step_handler_kwargs__' in data:
+                        del data['__step_handler_kwargs__']
+                    if '__step_handler_args__' in data:
+                        del data['__step_handler_args__']
+                await self.set_state(msg.from_user.id, None, msg.chat.id)
+                return await handler(*args, *step_args, **kwargs, **step_kwargs)
+            self.message_handlers.insert(0,{
+                'function': wrapper,
+                'pass_bot': False,
+                'filters': {'state': signature}
+            })
+            wrapper.signature = signature
+            return wrapper
+        return decorator
+    
+    async def register_next_step_handler(self, user_id: int, chat_id: int, callback: Callable, *args, **kwargs) -> None:
+        """
+        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.
+        :type message: :class:`telebot.types.Message`
+
+        :param callback: The callback function which next new message arrives.
+        :type callback: :obj:`Callable[[telebot.types.Message], None]`
+
+        :param args: Args to pass in callback func
+
+        :param kwargs: Args to pass in callback func
+
+        :return: None
+        """
+
+        if not hasattr(callback,'signature'):
+            raise ValueError("Do not forget to add @bot.step_handler() before function")
+        await self.set_state(user_id, callback.signature, chat_id)
+        await self.add_data(user_id, chat_id, __step_handler_args__=args, __step_handler_kwargs__=kwargs )
+
+
+    def callback_query_handler(self, func=None,**kwargs):
+        if func is None: func=lambda call:True
+        return super().callback_query_handler(func,**kwargs)
+            
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index e44c03a..d991d6e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,8 +5,8 @@
 # You can also run `make switch-to-poetry` to use the poetry package manager.
 
 i18nice==0.15.5
-#pyTelegramBotAPI==4.21.0
-git+https://github.com/hiddify-com/pyTelegramBotAPI.git@async_register_next_step_handler#egg=pyTelegramBotAPI
+pyTelegramBotAPI==4.21.0
+#git+https://github.com/hiddify-com/pyTelegramBotAPI.git@async_register_next_step_handler#egg=pyTelegramBotAPI
 aiohttp
 asyncio
 watchdog