From 12d0ae71dff2606bf3b2337bd5b85e308b04110f Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Thu, 26 Nov 2020 04:41:54 -0500
Subject: [PATCH] Add a way to check if the websocket is rate limited.

This is mainly for low level decision making for utilities that need
to know whether to fetch a member by HTTP or to query through the
websocket.

The library already does this trick in some places so it's only fair
that end users possess the same ability as well.
---
 discord/client.py | 12 ++++++++++++
 discord/shard.py  | 23 +++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/discord/client.py b/discord/client.py
index 1ddb09d6b..a5bb9c002 100644
--- a/discord/client.py
+++ b/discord/client.py
@@ -282,6 +282,18 @@ class Client:
         ws = self.ws
         return float('nan') if not ws else ws.latency
 
+    def is_ws_ratelimited(self):
+        """:class:`bool`: Whether the websocket is currently rate limited.
+
+        This can be useful to know when deciding whether you should query members
+        using HTTP or via the gateway.
+
+        .. versionadded:: 1.6
+        """
+        if self.ws:
+            return self.ws.is_ratelimited()
+        return False
+
     @property
     def user(self):
         """Optional[:class:`.ClientUser`]: Represents the connected client. ``None`` if not logged in."""
diff --git a/discord/shard.py b/discord/shard.py
index 6985d7970..ebc27b0da 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -258,6 +258,16 @@ class ShardInfo:
         """:class:`float`: Measures latency between a HEARTBEAT and a HEARTBEAT_ACK in seconds for this shard."""
         return self._parent.ws.latency
 
+    def is_ws_ratelimited(self):
+        """:class:`bool`: Whether the websocket is currently rate limited.
+
+        This can be useful to know when deciding whether you should query members
+        using HTTP or via the gateway.
+
+        .. versionadded:: 1.6
+        """
+        return self._parent.ws.is_ratelimited()
+
 class AutoShardedClient(Client):
     """A client similar to :class:`Client` except it handles the complications
     of sharding for the user into a more manageable and transparent single
@@ -519,3 +529,16 @@ class AutoShardedClient(Client):
 
             me.activities = activities
             me.status = status_enum
+
+    def is_ws_ratelimited(self):
+        """:class:`bool`: Whether the websocket is currently rate limited.
+
+        This can be useful to know when deciding whether you should query members
+        using HTTP or via the gateway.
+
+        This implementation checks if any of the shards are rate limited.
+        For more granular control, consider :meth:`ShardInfo.is_ws_ratelimited`.
+
+        .. versionadded:: 1.6
+        """
+        return any(shard.ws.is_ratelimited() for shard in self.__shards.values())