From 394b514cc927df61828017aa2078ff7dad6821e0 Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Sat, 11 Apr 2020 17:28:03 -0400
Subject: Add before_identify_hook to have finer control over IDENTIFY syncing

---
 discord/client.py | 43 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 37 insertions(+), 6 deletions(-)

(limited to 'discord/client.py')

diff --git a/discord/client.py b/discord/client.py
index 85931569..068e3d65 100644
--- a/discord/client.py
+++ b/discord/client.py
@@ -223,8 +223,12 @@ class Client:
             'ready': self._handle_ready
         }
 
+        self._hooks = {
+            'before_identify': self._call_before_identify_hook
+        }
+
         self._connection = ConnectionState(dispatch=self.dispatch, handlers=self._handlers,
-                                           syncer=self._syncer, http=self.http, loop=self.loop, **options)
+                                           hooks=self._hooks, syncer=self._syncer, http=self.http, loop=self.loop, **options)
 
         self._connection.shard_count = self.shard_count
         self._closed = False
@@ -394,6 +398,36 @@ class Client:
 
         await self._connection.request_offline_members(guilds)
 
+    # hooks
+
+    async def _call_before_identify_hook(self, shard_id, *, initial=False):
+        # This hook is an internal hook that actually calls the public one.
+        # It allows the library to have its own hook without stepping on the
+        # toes of those who need to override their own hook.
+        await self.before_identify_hook(shard_id, initial=initial)
+
+    async def before_identify_hook(self, shard_id, *, initial=False):
+        """|coro|
+
+        A hook that is called before IDENTIFYing a session. This is useful
+        if you wish to have more control over the synchronization of multiple
+        IDENTIFYing clients.
+
+        The default implementation sleeps for 5 seconds.
+
+        .. versionadded:: 1.4
+
+        Parameters
+        ------------
+        shard_id: :class:`int`
+            The shard ID that requested being IDENTIFY'd
+        initial: :class:`bool`
+            Whether this IDENTIFY is the first initial IDENTIFY.
+        """
+
+        if not initial:
+            await asyncio.sleep(5.0)
+
     # login state management
 
     async def login(self, token, *, bot=True):
@@ -447,7 +481,7 @@ class Client:
         await self.close()
 
     async def _connect(self):
-        coro = DiscordWebSocket.from_client(self, shard_id=self.shard_id)
+        coro = DiscordWebSocket.from_client(self, initial=True, shard_id=self.shard_id)
         self.ws = await asyncio.wait_for(coro, timeout=180.0)
         while True:
             try:
@@ -455,11 +489,8 @@ class Client:
             except ReconnectWebSocket as e:
                 log.info('Got a request to %s the websocket.', e.op)
                 self.dispatch('disconnect')
-                if not e.resume:
-                    await asyncio.sleep(5.0)
-
                 coro = DiscordWebSocket.from_client(self, shard_id=self.shard_id, session=self.ws.session_id,
-                                                    sequence=self.ws.sequence, resume=e.resume)
+                                                          sequence=self.ws.sequence, resume=e.resume)
                 self.ws = await asyncio.wait_for(coro, timeout=180.0)
 
     async def connect(self, *, reconnect=True):
-- 
cgit v1.2.3