Add an automatic sleep mechanism for flood wait exceptions

pyrogram/client/client.py

@@ -154,6 +154,12 @@ class Client(Methods, BaseClient):
             download_media, ...) are less prone to throw FloodWait exceptions.
             Only available for users, bots will ignore this parameter.
             Defaults to False (normal session).
+
+        sleep_threshold (``int``, *optional*):
+            Set a sleep threshold for flood wait exceptions happening globally in this client instance, below which any
+            request that raises a flood wait will be automatically invoked again after sleeping for the required amount
+            of time. Flood wait exceptions requiring higher waiting times will be raised.
+            Defaults to 60 (seconds).
     """
 
     def __init__(
@@ -178,7 +184,8 @@ class Client(Methods, BaseClient):
         config_file: str = BaseClient.CONFIG_FILE,
         plugins: dict = None,
         no_updates: bool = None,
-        takeout: bool = None
+        takeout: bool = None,
+        sleep_threshold: int = 60
     ):
         super().__init__()
 
@@ -204,6 +211,7 @@ class Client(Methods, BaseClient):
         self.plugins = plugins
         self.no_updates = no_updates
         self.takeout = takeout
+        self.sleep_threshold = sleep_threshold
 
         if isinstance(session_name, str):
             if session_name == ":memory:" or len(session_name) >= MemoryStorage.SESSION_STRING_SIZE:
@@ -1401,13 +1409,31 @@ class Client(Methods, BaseClient):
         if not self.is_connected:
             raise ConnectionError("Client has not been started yet")
 
+        # Some raw methods that expect a query as argument are used here.
+        # Keep the original request query because is needed.
+        unwrapped_data = data
+
         if self.no_updates:
             data = functions.InvokeWithoutUpdates(query=data)
 
         if self.takeout_id:
             data = functions.InvokeWithTakeout(takeout_id=self.takeout_id, query=data)
 
+        while True:
+            try:
                 r = self.session.send(data, retries, timeout)
+            except FloodWait as e:
+                amount = e.x
+
+                if amount > self.sleep_threshold:
+                    raise
+
+                log.warning('[{}] Sleeping for {}s (required by "{}")'.format(
+                    self.session_name, amount, ".".join(unwrapped_data.QUALNAME.split(".")[1:])))
+
+                time.sleep(amount)
+            else:
+                break
 
         self.fetch_peers(getattr(r, "users", []))
         self.fetch_peers(getattr(r, "chats", []))

pyrogram/client/methods/chats/get_chat_members.py

@@ -136,8 +136,6 @@ class GetChatMembers(BaseClient):
             else:
                 raise ValueError("Invalid filter \"{}\"".format(filter))
 
-            while True:
-                try:
             r = self.send(
                 functions.channels.GetParticipants(
                     channel=peer,
@@ -152,8 +150,5 @@ class GetChatMembers(BaseClient):
             users = {i.id: i for i in r.users}
 
             return pyrogram.List(pyrogram.ChatMember._parse(self, member, users) for member in members)
-                except FloodWait as e:
-                    log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x))
-                    time.sleep(e.x)
         else:
             raise ValueError("The chat_id \"{}\" belongs to a user".format(chat_id))

pyrogram/client/methods/chats/get_dialogs.py

@@ -66,8 +66,6 @@ class GetDialogs(BaseClient):
                 app.get_dialogs(pinned_only=True)
         """
 
-        while True:
-            try:
         if pinned_only:
             r = self.send(functions.messages.GetPinnedDialogs(folder_id=0))
         else:
@@ -81,11 +79,6 @@ class GetDialogs(BaseClient):
                     exclude_pinned=True
                 )
             )
-            except FloodWait as e:
-                log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x))
-                time.sleep(e.x)
-            else:
-                break
 
         users = {i.id: i for i in r.users}
         chats = {i.id: i for i in r.chats}

pyrogram/client/methods/messages/get_history.py

@@ -85,8 +85,6 @@ class GetHistory(BaseClient):
 
         offset_id = offset_id or (1 if reverse else 0)
 
-        while True:
-            try:
         messages = utils.parse_messages(
             self,
             self.send(
@@ -102,11 +100,6 @@ class GetHistory(BaseClient):
                 )
             )
         )
-            except FloodWait as e:
-                log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x))
-                time.sleep(e.x)
-            else:
-                break
 
         if reverse:
             messages.reverse()

pyrogram/client/methods/messages/get_messages.py

@@ -112,14 +112,7 @@ class GetMessages(BaseClient):
         else:
             rpc = functions.messages.GetMessages(id=ids)
 
-        while True:
-            try:
         r = self.send(rpc)
-            except FloodWait as e:
-                log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x))
-                time.sleep(e.x)
-            else:
-                break
 
         messages = utils.parse_messages(self, r, replies=replies)
 

pyrogram/client/methods/messages/send_media_group.py

@@ -78,8 +78,6 @@ class SendMediaGroup(BaseClient):
         for i in media:
             if isinstance(i, pyrogram.InputMediaPhoto):
                 if os.path.exists(i.media):
-                    while True:
-                        try:
                     media = self.send(
                         functions.messages.UploadMedia(
                             peer=self.resolve_peer(chat_id),
@@ -88,11 +86,6 @@ class SendMediaGroup(BaseClient):
                             )
                         )
                     )
-                        except FloodWait as e:
-                            log.warning("Sleeping for {}s".format(e.x))
-                            time.sleep(e.x)
-                        else:
-                            break
 
                     media = types.InputMediaPhoto(
                         id=types.InputPhoto(
@@ -122,8 +115,6 @@ class SendMediaGroup(BaseClient):
                     media = utils.get_input_media_from_file_id(i.media, i.file_ref, 2)
             elif isinstance(i, pyrogram.InputMediaVideo):
                 if os.path.exists(i.media):
-                    while True:
-                        try:
                     media = self.send(
                         functions.messages.UploadMedia(
                             peer=self.resolve_peer(chat_id),
@@ -143,11 +134,6 @@ class SendMediaGroup(BaseClient):
                             )
                         )
                     )
-                        except FloodWait as e:
-                            log.warning("Sleeping for {}s".format(e.x))
-                            time.sleep(e.x)
-                        else:
-                            break
 
                     media = types.InputMediaDocument(
                         id=types.InputDocument(
@@ -184,8 +170,6 @@ class SendMediaGroup(BaseClient):
                 )
             )
 
-        while True:
-            try:
         r = self.send(
             functions.messages.SendMultiMedia(
                 peer=self.resolve_peer(chat_id),
@@ -194,11 +178,6 @@ class SendMediaGroup(BaseClient):
                 reply_to_msg_id=reply_to_message_id
             )
         )
-            except FloodWait as e:
-                log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x))
-                time.sleep(e.x)
-            else:
-                break
 
         return utils.parse_messages(
             self,