Skip to content

Commit 31f39a0

Browse files
committed
Make plugin callback functions return the function itself when decorated
1 parent b4f0f41 commit 31f39a0

9 files changed

Lines changed: 79 additions & 108 deletions

File tree

pyrogram/client/client.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1106,7 +1106,7 @@ def load_plugins(self):
11061106
for name in vars(module).keys():
11071107
# noinspection PyBroadException
11081108
try:
1109-
handler, group = getattr(module, name)
1109+
handler, group = getattr(module, name).pyrogram_plugin
11101110

11111111
if isinstance(handler, Handler) and isinstance(group, int):
11121112
self.add_handler(handler, group)
@@ -1141,7 +1141,7 @@ def load_plugins(self):
11411141
for name in handlers:
11421142
# noinspection PyBroadException
11431143
try:
1144-
handler, group = getattr(module, name)
1144+
handler, group = getattr(module, name).pyrogram_plugin
11451145

11461146
if isinstance(handler, Handler) and isinstance(group, int):
11471147
self.add_handler(handler, group)
@@ -1179,7 +1179,7 @@ def load_plugins(self):
11791179
for name in handlers:
11801180
# noinspection PyBroadException
11811181
try:
1182-
handler, group = getattr(module, name)
1182+
handler, group = getattr(module, name).pyrogram_plugin
11831183

11841184
if isinstance(handler, Handler) and isinstance(group, int):
11851185
self.remove_handler(handler, group)

pyrogram/client/methods/decorators/on_callback_query.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -44,18 +43,15 @@ def on_callback_query(
4443
The group identifier, defaults to 0.
4544
"""
4645

47-
def decorator(func: callable) -> Tuple[Handler, int]:
48-
if isinstance(func, tuple):
49-
func = func[0].callback
46+
def decorator(func: Callable) -> Callable:
47+
if isinstance(self, pyrogram.Client):
48+
self.add_handler(pyrogram.CallbackQueryHandler(func, filters), group)
49+
elif isinstance(self, Filter) or self is None:
50+
func.pyrogram_plugin = (
51+
pyrogram.CallbackQueryHandler(func, self),
52+
group if filters is None else filters
53+
)
5054

51-
handler = pyrogram.CallbackQueryHandler(func, filters)
52-
53-
if isinstance(self, Filter):
54-
return pyrogram.CallbackQueryHandler(func, self), group if filters is None else filters
55-
56-
if self is not None:
57-
self.add_handler(handler, group)
58-
59-
return handler, group
55+
return func
6056

6157
return decorator

pyrogram/client/methods/decorators/on_deleted_messages.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -44,18 +43,15 @@ def on_deleted_messages(
4443
The group identifier, defaults to 0.
4544
"""
4645

47-
def decorator(func: callable) -> Tuple[Handler, int]:
48-
if isinstance(func, tuple):
49-
func = func[0].callback
46+
def decorator(func: Callable) -> Callable:
47+
if isinstance(self, pyrogram.Client):
48+
self.add_handler(pyrogram.DeletedMessagesHandler(func, filters), group)
49+
elif isinstance(self, Filter) or self is None:
50+
func.pyrogram_plugin = (
51+
pyrogram.DeletedMessagesHandler(func, self),
52+
group if filters is None else filters
53+
)
5054

51-
handler = pyrogram.DeletedMessagesHandler(func, filters)
52-
53-
if isinstance(self, Filter):
54-
return pyrogram.DeletedMessagesHandler(func, self), group if filters is None else filters
55-
56-
if self is not None:
57-
self.add_handler(handler, group)
58-
59-
return handler, group
55+
return func
6056

6157
return decorator

pyrogram/client/methods/decorators/on_disconnect.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,9 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19+
from typing import Callable
20+
1921
import pyrogram
20-
from pyrogram.client.handlers.handler import Handler
2122
from ...ext import BaseClient
2223

2324

@@ -28,12 +29,10 @@ def on_disconnect(self=None) -> callable:
2829
This does the same thing as :meth:`~pyrogram.Client.add_handler` using the :obj:`~pyrogram.DisconnectHandler`.
2930
"""
3031

31-
def decorator(func: callable) -> Handler:
32-
handler = pyrogram.DisconnectHandler(func)
33-
34-
if self is not None:
35-
self.add_handler(handler)
32+
def decorator(func: Callable) -> Callable:
33+
if isinstance(self, pyrogram.Client):
34+
self.add_handler(pyrogram.DisconnectHandler(func))
3635

37-
return handler
36+
return func
3837

3938
return decorator

pyrogram/client/methods/decorators/on_inline_query.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -43,18 +42,15 @@ def on_inline_query(
4342
The group identifier, defaults to 0.
4443
"""
4544

46-
def decorator(func: callable) -> Tuple[Handler, int]:
47-
if isinstance(func, tuple):
48-
func = func[0].callback
45+
def decorator(func: Callable) -> Callable:
46+
if isinstance(self, pyrogram.Client):
47+
self.add_handler(pyrogram.InlineQueryHandler(func, filters), group)
48+
elif isinstance(self, Filter) or self is None:
49+
func.pyrogram_plugin = (
50+
pyrogram.InlineQueryHandler(func, self),
51+
group if filters is None else filters
52+
)
4953

50-
handler = pyrogram.InlineQueryHandler(func, filters)
51-
52-
if isinstance(self, Filter):
53-
return pyrogram.InlineQueryHandler(func, self), group if filters is None else filters
54-
55-
if self is not None:
56-
self.add_handler(handler, group)
57-
58-
return handler, group
54+
return func
5955

6056
return decorator

pyrogram/client/methods/decorators/on_message.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -43,18 +42,15 @@ def on_message(
4342
The group identifier, defaults to 0.
4443
"""
4544

46-
def decorator(func: callable) -> Tuple[Handler, int]:
47-
if isinstance(func, tuple):
48-
func = func[0].callback
45+
def decorator(func: Callable) -> Callable:
46+
if isinstance(self, pyrogram.Client):
47+
self.add_handler(pyrogram.MessageHandler(func, filters), group)
48+
elif isinstance(self, Filter) or self is None:
49+
func.pyrogram_plugin = (
50+
pyrogram.MessageHandler(func, self),
51+
group if filters is None else filters
52+
)
4953

50-
handler = pyrogram.MessageHandler(func, filters)
51-
52-
if isinstance(self, Filter):
53-
return pyrogram.MessageHandler(func, self), group if filters is None else filters
54-
55-
if self is not None:
56-
self.add_handler(handler, group)
57-
58-
return handler, group
54+
return func
5955

6056
return decorator

pyrogram/client/methods/decorators/on_poll.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -43,18 +42,15 @@ def on_poll(
4342
The group identifier, defaults to 0.
4443
"""
4544

46-
def decorator(func: callable) -> Tuple[Handler, int]:
47-
if isinstance(func, tuple):
48-
func = func[0].callback
45+
def decorator(func: Callable) -> Callable:
46+
if isinstance(self, pyrogram.Client):
47+
self.add_handler(pyrogram.PollHandler(func, filters), group)
48+
elif isinstance(self, Filter) or self is None:
49+
func.pyrogram_plugin = (
50+
pyrogram.PollHandler(func, self),
51+
group if filters is None else filters
52+
)
4953

50-
handler = pyrogram.PollHandler(func, filters)
51-
52-
if isinstance(self, Filter):
53-
return pyrogram.PollHandler(func, self), group if filters is None else filters
54-
55-
if self is not None:
56-
self.add_handler(handler, group)
57-
58-
return handler, group
54+
return func
5955

6056
return decorator

pyrogram/client/methods/decorators/on_raw_update.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,9 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
22-
from pyrogram.client.handlers.handler import Handler
2322
from ...ext import BaseClient
2423

2524

@@ -37,18 +36,15 @@ def on_raw_update(
3736
The group identifier, defaults to 0.
3837
"""
3938

40-
def decorator(func: callable) -> Tuple[Handler, int]:
41-
if isinstance(func, tuple):
42-
func = func[0].callback
39+
def decorator(func: Callable) -> Callable:
40+
if isinstance(self, pyrogram.Client):
41+
self.add_handler(pyrogram.RawUpdateHandler(func), group)
42+
else:
43+
func.pyrogram_plugin = (
44+
pyrogram.RawUpdateHandler(func),
45+
group if self is None else group
46+
)
4347

44-
handler = pyrogram.RawUpdateHandler(func)
45-
46-
if isinstance(self, int):
47-
return handler, group if self is None else group
48-
49-
if self is not None:
50-
self.add_handler(handler, group)
51-
52-
return handler, group
48+
return func
5349

5450
return decorator

pyrogram/client/methods/decorators/on_user_status.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19-
from typing import Tuple
19+
from typing import Callable
2020

2121
import pyrogram
2222
from pyrogram.client.filters.filter import Filter
23-
from pyrogram.client.handlers.handler import Handler
2423
from ...ext import BaseClient
2524

2625

@@ -41,18 +40,15 @@ def on_user_status(
4140
The group identifier, defaults to 0.
4241
"""
4342

44-
def decorator(func: callable) -> Tuple[Handler, int]:
45-
if isinstance(func, tuple):
46-
func = func[0].callback
43+
def decorator(func: Callable) -> Callable:
44+
if isinstance(self, pyrogram.Client):
45+
self.add_handler(pyrogram.UserStatusHandler(func, filters), group)
46+
elif isinstance(self, Filter) or self is None:
47+
func.pyrogram_plugin = (
48+
pyrogram.UserStatusHandler(func, self),
49+
group if filters is None else filters
50+
)
4751

48-
handler = pyrogram.UserStatusHandler(func, filters)
49-
50-
if isinstance(self, Filter):
51-
return pyrogram.UserStatusHandler(func, self), group if filters is None else filters
52-
53-
if self is not None:
54-
self.add_handler(handler, group)
55-
56-
return handler, group
52+
return func
5753

5854
return decorator

0 commit comments

Comments
 (0)