diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 0e70fb5b..90facb8e 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -81,6 +81,12 @@ class Client: self.config = None self.session = None + self.update_handler = None + + # TODO: Better update handler + def set_update_handler(self, callback): + self.update_handler = callback + def send(self, data: Object): return self.session.send(data) @@ -249,6 +255,7 @@ class Client: self.load_session(self.session_name) self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id) + self.session.update_handler = self.update_handler terms = self.session.start() diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index 1a46ce56..584bd29c 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -96,6 +96,8 @@ class Session: self.is_connected = Event() + self.update_handler = None + self.total_connections = 0 self.total_messages = 0 self.total_bytes = 0 @@ -246,19 +248,22 @@ class Session: self.pending_acks.add(i.body.answer_msg_id) continue + if isinstance(i.body, types.NewSessionCreated): + continue + msg_id = None if isinstance(i.body, (types.BadMsgNotification, types.BadServerSalt)): msg_id = i.body.bad_msg_id - elif isinstance(i.body, types.RpcResult): msg_id = i.body.req_msg_id - elif isinstance(i.body, types.Pong): msg_id = i.body.msg_id - elif isinstance(i.body, core.FutureSalts): msg_id = i.body.req_msg_id + else: + if self.update_handler: + self.update_handler(i.body) if msg_id in self.results: self.results[msg_id].value = getattr(i.body, "result", i.body)