Skip to content

Commit 9001ccd

Browse files
delivrancedelivrance
committed
Add DisconnectHandler
1 parent 8a5743e commit 9001ccd

9 files changed

Lines changed: 80 additions & 5 deletions

File tree

‎pyrogram/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,5 +39,5 @@
3939
from .client import (
4040
Client, ChatAction, ParseMode, Emoji,
4141
MessageHandler, CallbackQueryHandler, RawUpdateHandler,
42-
Filters
42+
DisconnectHandler, Filters
4343
)

‎pyrogram/client/__init__.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,7 @@
1919
from .client import Client
2020
from .ext import BaseClient, ChatAction, Emoji, ParseMode
2121
from .filters import Filters
22-
from .handlers import MessageHandler, CallbackQueryHandler, RawUpdateHandler
22+
from .handlers import (
23+
MessageHandler, CallbackQueryHandler,
24+
RawUpdateHandler, DisconnectHandler
25+
)

‎pyrogram/client/client.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded,
4444
PasswordHashInvalid, FloodWait, PeerIdInvalid, FirstnameInvalid, PhoneNumberBanned,
4545
VolumeLocNotFound, UserMigrate, FileIdInvalid)
46+
from pyrogram.client.handlers import DisconnectHandler
4647
from pyrogram.crypto import AES
4748
from pyrogram.session import Auth, Session
4849
from .dispatcher import Dispatcher
@@ -290,7 +291,10 @@ def add_handler(self, handler, group: int = 0):
290291
Returns:
291292
A tuple of (handler, group)
292293
"""
293-
self.dispatcher.add_handler(handler, group)
294+
if isinstance(handler, DisconnectHandler):
295+
self.disconnect_handler = handler.callback
296+
else:
297+
self.dispatcher.add_handler(handler, group)
294298

295299
return handler, group
296300

@@ -308,7 +312,10 @@ def remove_handler(self, handler, group: int = 0):
308312
group (``int``, *optional*):
309313
The group identifier, defaults to 0.
310314
"""
311-
self.dispatcher.remove_handler(handler, group)
315+
if isinstance(handler, DisconnectHandler):
316+
self.disconnect_handler = None
317+
else:
318+
self.dispatcher.remove_handler(handler, group)
312319

313320
def authorize_bot(self):
314321
try:

‎pyrogram/client/ext/base_client.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ def __init__(self):
7575
self.download_queue = Queue()
7676
self.download_workers_list = []
7777

78+
self.disconnect_handler = None
79+
7880
def send(self, data: Object):
7981
pass
8082

‎pyrogram/client/handlers/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,5 +17,6 @@
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

1919
from .callback_query_handler import CallbackQueryHandler
20+
from .disconnect_handler import DisconnectHandler
2021
from .message_handler import MessageHandler
2122
from .raw_update_handler import RawUpdateHandler

‎pyrogram/client/handlers/disconnect_handler.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Pyrogram - Telegram MTProto API Client Library for Python
2+
# Copyright (C) 2017-2018 Dan Tès <https://github.com/delivrance>
3+
#
4+
# This file is part of Pyrogram.
5+
#
6+
# Pyrogram is free software: you can redistribute it and/or modify
7+
# it under the terms of the GNU Lesser General Public License as published
8+
# by the Free Software Foundation, either version 3 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# Pyrogram is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU Lesser General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU Lesser General Public License
17+
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
18+
19+
from .handler import Handler
20+
21+
22+
class DisconnectHandler(Handler):
23+
# TODO: Documentation
24+
def __init__(self, callback: callable):
25+
super().__init__(callback)

‎pyrogram/client/methods/decorators/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,10 @@
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

1919
from .on_callback_query import OnCallbackQuery
20+
from .on_disconnect import OnDisconnect
2021
from .on_message import OnMessage
2122
from .on_raw_update import OnRawUpdate
2223

2324

24-
class Decorators(OnMessage, OnCallbackQuery, OnRawUpdate):
25+
class Decorators(OnMessage, OnCallbackQuery, OnRawUpdate, OnDisconnect):
2526
pass

‎pyrogram/client/methods/decorators/on_disconnect.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Pyrogram - Telegram MTProto API Client Library for Python
2+
# Copyright (C) 2017-2018 Dan Tès <https://github.com/delivrance>
3+
#
4+
# This file is part of Pyrogram.
5+
#
6+
# Pyrogram is free software: you can redistribute it and/or modify
7+
# it under the terms of the GNU Lesser General Public License as published
8+
# by the Free Software Foundation, either version 3 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# Pyrogram is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU Lesser General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU Lesser General Public License
17+
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
18+
19+
import pyrogram
20+
from ...ext import BaseClient
21+
22+
23+
class OnDisconnect(BaseClient):
24+
def on_disconnect(self):
25+
# TODO: Documentation
26+
def decorator(func):
27+
self.add_handler(pyrogram.DisconnectHandler(func))
28+
return func
29+
30+
return decorator

‎pyrogram/session/session.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,12 @@ def stop(self):
207207
for i in self.results.values():
208208
i.event.set()
209209

210+
if self.client and callable(self.client.disconnect_handler):
211+
try:
212+
self.client.disconnect_handler(self.client)
213+
except Exception as e:
214+
log.error(e, exc_info=True)
215+
210216
log.debug("Session stopped")
211217

212218
def restart(self):

0 commit comments

Comments
 (0)