From 9001ccd11ff2875fb760be1e1ced1afa46df1c60 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 23 May 2018 14:27:17 +0200
Subject: [PATCH] Add DisconnectHandler

---
 pyrogram/__init__.py                          |  2 +-
 pyrogram/client/__init__.py                   |  5 +++-
 pyrogram/client/client.py                     | 11 +++++--
 pyrogram/client/ext/base_client.py            |  2 ++
 pyrogram/client/handlers/__init__.py          |  1 +
 .../client/handlers/disconnect_handler.py     | 25 ++++++++++++++++
 .../client/methods/decorators/__init__.py     |  3 +-
 .../methods/decorators/on_disconnect.py       | 30 +++++++++++++++++++
 pyrogram/session/session.py                   |  6 ++++
 9 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 pyrogram/client/handlers/disconnect_handler.py
 create mode 100644 pyrogram/client/methods/decorators/on_disconnect.py

diff --git a/pyrogram/__init__.py b/pyrogram/__init__.py
index 636e52bf..5581b969 100644
--- a/pyrogram/__init__.py
+++ b/pyrogram/__init__.py
@@ -39,5 +39,5 @@ from .client.types.reply_markup import (
 from .client import (
     Client, ChatAction, ParseMode, Emoji,
     MessageHandler, CallbackQueryHandler, RawUpdateHandler,
-    Filters
+    DisconnectHandler, Filters
 )
diff --git a/pyrogram/client/__init__.py b/pyrogram/client/__init__.py
index aac1b9e0..21bde533 100644
--- a/pyrogram/client/__init__.py
+++ b/pyrogram/client/__init__.py
@@ -19,4 +19,7 @@
 from .client import Client
 from .ext import BaseClient, ChatAction, Emoji, ParseMode
 from .filters import Filters
-from .handlers import MessageHandler, CallbackQueryHandler, RawUpdateHandler
+from .handlers import (
+    MessageHandler, CallbackQueryHandler,
+    RawUpdateHandler, DisconnectHandler
+)
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index 8721f925..75bb082f 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -43,6 +43,7 @@ from pyrogram.api.errors import (
     PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded,
     PasswordHashInvalid, FloodWait, PeerIdInvalid, FirstnameInvalid, PhoneNumberBanned,
     VolumeLocNotFound, UserMigrate, FileIdInvalid)
+from pyrogram.client.handlers import DisconnectHandler
 from pyrogram.crypto import AES
 from pyrogram.session import Auth, Session
 from .dispatcher import Dispatcher
@@ -290,7 +291,10 @@ class Client(Methods, BaseClient):
         Returns:
             A tuple of (handler, group)
         """
-        self.dispatcher.add_handler(handler, group)
+        if isinstance(handler, DisconnectHandler):
+            self.disconnect_handler = handler.callback
+        else:
+            self.dispatcher.add_handler(handler, group)
 
         return handler, group
 
@@ -308,7 +312,10 @@ class Client(Methods, BaseClient):
             group (``int``, *optional*):
                 The group identifier, defaults to 0.
         """
-        self.dispatcher.remove_handler(handler, group)
+        if isinstance(handler, DisconnectHandler):
+            self.disconnect_handler = None
+        else:
+            self.dispatcher.remove_handler(handler, group)
 
     def authorize_bot(self):
         try:
diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py
index ce3d5591..9c0fb26b 100644
--- a/pyrogram/client/ext/base_client.py
+++ b/pyrogram/client/ext/base_client.py
@@ -75,6 +75,8 @@ class BaseClient:
         self.download_queue = Queue()
         self.download_workers_list = []
 
+        self.disconnect_handler = None
+
     def send(self, data: Object):
         pass
 
diff --git a/pyrogram/client/handlers/__init__.py b/pyrogram/client/handlers/__init__.py
index 52fe40b3..38a8aeeb 100644
--- a/pyrogram/client/handlers/__init__.py
+++ b/pyrogram/client/handlers/__init__.py
@@ -17,5 +17,6 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from .callback_query_handler import CallbackQueryHandler
+from .disconnect_handler import DisconnectHandler
 from .message_handler import MessageHandler
 from .raw_update_handler import RawUpdateHandler
diff --git a/pyrogram/client/handlers/disconnect_handler.py b/pyrogram/client/handlers/disconnect_handler.py
new file mode 100644
index 00000000..67d0ae02
--- /dev/null
+++ b/pyrogram/client/handlers/disconnect_handler.py
@@ -0,0 +1,25 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2018 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# 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 .handler import Handler
+
+
+class DisconnectHandler(Handler):
+    # TODO: Documentation
+    def __init__(self, callback: callable):
+        super().__init__(callback)
diff --git a/pyrogram/client/methods/decorators/__init__.py b/pyrogram/client/methods/decorators/__init__.py
index 89645906..6ece61ff 100644
--- a/pyrogram/client/methods/decorators/__init__.py
+++ b/pyrogram/client/methods/decorators/__init__.py
@@ -17,9 +17,10 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from .on_callback_query import OnCallbackQuery
+from .on_disconnect import OnDisconnect
 from .on_message import OnMessage
 from .on_raw_update import OnRawUpdate
 
 
-class Decorators(OnMessage, OnCallbackQuery, OnRawUpdate):
+class Decorators(OnMessage, OnCallbackQuery, OnRawUpdate, OnDisconnect):
     pass
diff --git a/pyrogram/client/methods/decorators/on_disconnect.py b/pyrogram/client/methods/decorators/on_disconnect.py
new file mode 100644
index 00000000..230a361c
--- /dev/null
+++ b/pyrogram/client/methods/decorators/on_disconnect.py
@@ -0,0 +1,30 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2018 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import pyrogram
+from ...ext import BaseClient
+
+
+class OnDisconnect(BaseClient):
+    def on_disconnect(self):
+        # TODO: Documentation
+        def decorator(func):
+            self.add_handler(pyrogram.DisconnectHandler(func))
+            return func
+
+        return decorator
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 2564b7d1..1da0fdd7 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -207,6 +207,12 @@ class Session:
         for i in self.results.values():
             i.event.set()
 
+        if self.client and callable(self.client.disconnect_handler):
+            try:
+                self.client.disconnect_handler(self.client)
+            except Exception as e:
+                log.error(e, exc_info=True)
+
         log.debug("Session stopped")
 
     def restart(self):