From 31f39a00ab602c9714c827833e93730b9466bcb2 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 23 Jun 2019 01:33:46 +0200
Subject: [PATCH] Make plugin callback functions return the function itself
 when decorated

---
 pyrogram/client/client.py                     |  6 ++---
 .../methods/decorators/on_callback_query.py   | 24 ++++++++-----------
 .../methods/decorators/on_deleted_messages.py | 24 ++++++++-----------
 .../methods/decorators/on_disconnect.py       | 13 +++++-----
 .../methods/decorators/on_inline_query.py     | 24 ++++++++-----------
 .../client/methods/decorators/on_message.py   | 24 ++++++++-----------
 pyrogram/client/methods/decorators/on_poll.py | 24 ++++++++-----------
 .../methods/decorators/on_raw_update.py       | 24 ++++++++-----------
 .../methods/decorators/on_user_status.py      | 24 ++++++++-----------
 9 files changed, 79 insertions(+), 108 deletions(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index e1bebeaf..fdfc1e68 100644
--- a/pyrogram/client/client.py
+++ b/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)
diff --git a/pyrogram/client/methods/decorators/on_callback_query.py b/pyrogram/client/methods/decorators/on_callback_query.py
index 1552bae7..1706d71a 100644
--- a/pyrogram/client/methods/decorators/on_callback_query.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_deleted_messages.py b/pyrogram/client/methods/decorators/on_deleted_messages.py
index 0d87ba5a..86dda587 100644
--- a/pyrogram/client/methods/decorators/on_deleted_messages.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_disconnect.py b/pyrogram/client/methods/decorators/on_disconnect.py
index 4a514a41..012abd38 100644
--- a/pyrogram/client/methods/decorators/on_disconnect.py
+++ b/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:
-            handler = pyrogram.DisconnectHandler(func)
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                self.add_handler(pyrogram.DisconnectHandler(func))
 
-            if self is not None:
-                self.add_handler(handler)
-
-            return handler
+            return func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_inline_query.py b/pyrogram/client/methods/decorators/on_inline_query.py
index adc65d25..d0f2925b 100644
--- a/pyrogram/client/methods/decorators/on_inline_query.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_message.py b/pyrogram/client/methods/decorators/on_message.py
index 758a6831..5640f22c 100644
--- a/pyrogram/client/methods/decorators/on_message.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_poll.py b/pyrogram/client/methods/decorators/on_poll.py
index 0ade42c0..24282f28 100644
--- a/pyrogram/client/methods/decorators/on_poll.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_raw_update.py b/pyrogram/client/methods/decorators/on_raw_update.py
index 7dff75fa..bbf40c8b 100644
--- a/pyrogram/client/methods/decorators/on_raw_update.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator
diff --git a/pyrogram/client/methods/decorators/on_user_status.py b/pyrogram/client/methods/decorators/on_user_status.py
index 09e037f7..81a83d02 100644
--- a/pyrogram/client/methods/decorators/on_user_status.py
+++ b/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]:
-            if isinstance(func, tuple):
-                func = func[0].callback
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                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)
-
-            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 func
 
         return decorator