Make plugin callback functions return the function itself when decorated

pyrogram/client/client.py

@@ -1106,7 +1106,7 @@ class Client(Methods, BaseClient):
                     for name in vars(module).keys():
                         # noinspection PyBroadException
                         try:
-                            handler, group = getattr(module, name)
+                            handler, group = getattr(module, name).pyrogram_plugin
 
                             if isinstance(handler, Handler) and isinstance(group, int):
                                 self.add_handler(handler, group)
@@ -1141,7 +1141,7 @@ class Client(Methods, BaseClient):
                     for name in handlers:
                         # noinspection PyBroadException
                         try:
-                            handler, group = getattr(module, name)
+                            handler, group = getattr(module, name).pyrogram_plugin
 
                             if isinstance(handler, Handler) and isinstance(group, int):
                                 self.add_handler(handler, group)
@@ -1179,7 +1179,7 @@ class Client(Methods, BaseClient):
                     for name in handlers:
                         # noinspection PyBroadException
                         try:
-                            handler, group = getattr(module, name)
+                            handler, group = getattr(module, name).pyrogram_plugin
 
                             if isinstance(handler, Handler) and isinstance(group, int):
                                 self.remove_handler(handler, group)

pyrogram/client/methods/decorators/on_callback_query.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -44,18 +43,15 @@ class OnCallbackQuery(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.CallbackQueryHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.CallbackQueryHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.CallbackQueryHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.CallbackQueryHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_deleted_messages.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -44,18 +43,15 @@ class OnDeletedMessages(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.DeletedMessagesHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.DeletedMessagesHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.DeletedMessagesHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.DeletedMessagesHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_disconnect.py

@@ -16,8 +16,9 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+from typing import Callable
+
 import pyrogram
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -28,12 +29,10 @@ class OnDisconnect(BaseClient):
         This does the same thing as :meth:`~pyrogram.Client.add_handler` using the :obj:`~pyrogram.DisconnectHandler`.
         """
 
-        def decorator(func: callable) -> Handler:
+        def decorator(func: Callable) -> Callable:
-            handler = pyrogram.DisconnectHandler(func)
+            if isinstance(self, pyrogram.Client):
+                self.add_handler(pyrogram.DisconnectHandler(func))
 
-            if self is not None:
+            return func
-                self.add_handler(handler)
-
-            return handler
 
         return decorator

pyrogram/client/methods/decorators/on_inline_query.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -43,18 +42,15 @@ class OnInlineQuery(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.InlineQueryHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.InlineQueryHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.InlineQueryHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.InlineQueryHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_message.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -43,18 +42,15 @@ class OnMessage(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.MessageHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.MessageHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.MessageHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.MessageHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_poll.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -43,18 +42,15 @@ class OnPoll(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.PollHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.PollHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.PollHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.PollHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_raw_update.py

@@ -16,10 +16,9 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -37,18 +36,15 @@ class OnRawUpdate(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.RawUpdateHandler(func), group)
+            else:
+                func.pyrogram_plugin = (
+                    pyrogram.RawUpdateHandler(func),
+                    group if self is None else group
+                )
 
-            handler = pyrogram.RawUpdateHandler(func)
+            return func
-
-            if isinstance(self, int):
-                return handler, group if self is None else group
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator

pyrogram/client/methods/decorators/on_user_status.py

@@ -16,11 +16,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Tuple
+from typing import Callable
 
 import pyrogram
 from pyrogram.client.filters.filter import Filter
-from pyrogram.client.handlers.handler import Handler
 from ...ext import BaseClient
 
 
@@ -41,18 +40,15 @@ class OnUserStatus(BaseClient):
                 The group identifier, defaults to 0.
         """
 
-        def decorator(func: callable) -> Tuple[Handler, int]:
+        def decorator(func: Callable) -> Callable:
-            if isinstance(func, tuple):
+            if isinstance(self, pyrogram.Client):
-                func = func[0].callback
+                self.add_handler(pyrogram.UserStatusHandler(func, filters), group)
+            elif isinstance(self, Filter) or self is None:
+                func.pyrogram_plugin = (
+                    pyrogram.UserStatusHandler(func, self),
+                    group if filters is None else filters
+                )
 
-            handler = pyrogram.UserStatusHandler(func, filters)
+            return func
-
-            if isinstance(self, Filter):
-                return pyrogram.UserStatusHandler(func, self), group if filters is None else filters
-
-            if self is not None:
-                self.add_handler(handler, group)
-
-            return handler, group
 
         return decorator