diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 3ef82ada..daaa8669 100644 --- a/pyrogram/client/client.py +++ b/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) - r = self.session.send(data, retries, timeout) + 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", [])) diff --git a/pyrogram/client/methods/chats/get_chat_members.py b/pyrogram/client/methods/chats/get_chat_members.py index 6be3408b..bf208944 100644 --- a/pyrogram/client/methods/chats/get_chat_members.py +++ b/pyrogram/client/methods/chats/get_chat_members.py @@ -136,24 +136,19 @@ class GetChatMembers(BaseClient): else: raise ValueError("Invalid filter \"{}\"".format(filter)) - while True: - try: - r = self.send( - functions.channels.GetParticipants( - channel=peer, - filter=filter, - offset=offset, - limit=limit, - hash=0 - ) - ) + r = self.send( + functions.channels.GetParticipants( + channel=peer, + filter=filter, + offset=offset, + limit=limit, + hash=0 + ) + ) - members = r.participants - users = {i.id: i for i in r.users} + members = r.participants + 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) + return pyrogram.List(pyrogram.ChatMember._parse(self, member, users) for member in members) else: raise ValueError("The chat_id \"{}\" belongs to a user".format(chat_id)) diff --git a/pyrogram/client/methods/chats/get_dialogs.py b/pyrogram/client/methods/chats/get_dialogs.py index ea2eec63..f5d5f442 100644 --- a/pyrogram/client/methods/chats/get_dialogs.py +++ b/pyrogram/client/methods/chats/get_dialogs.py @@ -66,26 +66,19 @@ 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: - r = self.send( - functions.messages.GetDialogs( - offset_date=offset_date, - offset_id=0, - offset_peer=types.InputPeerEmpty(), - limit=limit, - hash=0, - exclude_pinned=True - ) - ) - except FloodWait as e: - log.warning("[{}] Sleeping for {}s".format(self.session_name, e.x)) - time.sleep(e.x) - else: - break + if pinned_only: + r = self.send(functions.messages.GetPinnedDialogs(folder_id=0)) + else: + r = self.send( + functions.messages.GetDialogs( + offset_date=offset_date, + offset_id=0, + offset_peer=types.InputPeerEmpty(), + limit=limit, + hash=0, + exclude_pinned=True + ) + ) users = {i.id: i for i in r.users} chats = {i.id: i for i in r.chats} diff --git a/pyrogram/client/methods/messages/get_history.py b/pyrogram/client/methods/messages/get_history.py index 7741340c..79a1dec6 100644 --- a/pyrogram/client/methods/messages/get_history.py +++ b/pyrogram/client/methods/messages/get_history.py @@ -85,28 +85,21 @@ class GetHistory(BaseClient): offset_id = offset_id or (1 if reverse else 0) - while True: - try: - messages = utils.parse_messages( - self, - self.send( - functions.messages.GetHistory( - peer=self.resolve_peer(chat_id), - offset_id=offset_id, - offset_date=offset_date, - add_offset=offset * (-1 if reverse else 1) - (limit if reverse else 0), - limit=limit, - max_id=0, - min_id=0, - hash=0 - ) - ) + messages = utils.parse_messages( + self, + self.send( + functions.messages.GetHistory( + peer=self.resolve_peer(chat_id), + offset_id=offset_id, + offset_date=offset_date, + add_offset=offset * (-1 if reverse else 1) - (limit if reverse else 0), + limit=limit, + max_id=0, + min_id=0, + hash=0 ) - 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() diff --git a/pyrogram/client/methods/messages/get_messages.py b/pyrogram/client/methods/messages/get_messages.py index f3e53b94..e505b566 100644 --- a/pyrogram/client/methods/messages/get_messages.py +++ b/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 + r = self.send(rpc) messages = utils.parse_messages(self, r, replies=replies) diff --git a/pyrogram/client/methods/messages/send_media_group.py b/pyrogram/client/methods/messages/send_media_group.py index 725a8a84..8571ef4f 100644 --- a/pyrogram/client/methods/messages/send_media_group.py +++ b/pyrogram/client/methods/messages/send_media_group.py @@ -78,21 +78,14 @@ 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), - media=types.InputMediaUploadedPhoto( - file=self.save_file(i.media) - ) - ) + media = self.send( + functions.messages.UploadMedia( + peer=self.resolve_peer(chat_id), + media=types.InputMediaUploadedPhoto( + file=self.save_file(i.media) ) - 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,32 +115,25 @@ 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), - media=types.InputMediaUploadedDocument( - file=self.save_file(i.media), - thumb=None if i.thumb is None else self.save_file(i.thumb), - mime_type=self.guess_mime_type(i.media) or "video/mp4", - attributes=[ - types.DocumentAttributeVideo( - supports_streaming=i.supports_streaming or None, - duration=i.duration, - w=i.width, - h=i.height - ), - types.DocumentAttributeFilename(file_name=os.path.basename(i.media)) - ] - ) - ) + media = self.send( + functions.messages.UploadMedia( + peer=self.resolve_peer(chat_id), + media=types.InputMediaUploadedDocument( + file=self.save_file(i.media), + thumb=None if i.thumb is None else self.save_file(i.thumb), + mime_type=self.guess_mime_type(i.media) or "video/mp4", + attributes=[ + types.DocumentAttributeVideo( + supports_streaming=i.supports_streaming or None, + duration=i.duration, + w=i.width, + h=i.height + ), + types.DocumentAttributeFilename(file_name=os.path.basename(i.media)) + ] ) - 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,21 +170,14 @@ class SendMediaGroup(BaseClient): ) ) - while True: - try: - r = self.send( - functions.messages.SendMultiMedia( - peer=self.resolve_peer(chat_id), - multi_media=multi_media, - silent=disable_notification or None, - 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 + r = self.send( + functions.messages.SendMultiMedia( + peer=self.resolve_peer(chat_id), + multi_media=multi_media, + silent=disable_notification or None, + reply_to_msg_id=reply_to_message_id + ) + ) return utils.parse_messages( self,