From 128ab4b0b94851720aa30782546d22195fc39f93 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 7 May 2020 13:38:22 +0200
Subject: [PATCH] Move the automatic sleep mechanism down to Session

---
 pyrogram/client/client.py   | 22 ++--------------
 pyrogram/session/session.py | 52 +++++++++++++++++++++++++++----------
 2 files changed, 41 insertions(+), 33 deletions(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index daaa8669..53bc3799 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -185,7 +185,7 @@ class Client(Methods, BaseClient):
         plugins: dict = None,
         no_updates: bool = None,
         takeout: bool = None,
-        sleep_threshold: int = 60
+        sleep_threshold: int = Session.SLEEP_THRESHOLD
     ):
         super().__init__()
 
@@ -1409,31 +1409,13 @@ 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
+        r = self.session.send(data, retries, timeout, self.sleep_threshold)
 
         self.fetch_peers(getattr(r, "users", []))
         self.fetch_peers(getattr(r, "chats", []))
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index e2659871..3c34f578 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -33,7 +33,7 @@ from pyrogram.api.all import layer
 from pyrogram.api.core import Message, TLObject, MsgContainer, Long, FutureSalt, Int
 from pyrogram.connection import Connection
 from pyrogram.crypto import AES, KDF
-from pyrogram.errors import RPCError, InternalServerError, AuthKeyDuplicated
+from pyrogram.errors import RPCError, InternalServerError, AuthKeyDuplicated, FloodWait
 from .internals import MsgId, MsgFactory
 
 log = logging.getLogger(__name__)
@@ -50,6 +50,7 @@ class Session:
     NET_WORKERS = 1
     START_TIMEOUT = 1
     WAIT_TIMEOUT = 15
+    SLEEP_THRESHOLD = 60
     MAX_RETRIES = 5
     ACKS_THRESHOLD = 8
     PING_INTERVAL = 5
@@ -432,19 +433,44 @@ class Session:
             else:
                 return result
 
-    def send(self, data: TLObject, retries: int = MAX_RETRIES, timeout: float = WAIT_TIMEOUT):
+    def send(
+        self,
+        data: TLObject,
+        retries: int = MAX_RETRIES,
+        timeout: float = WAIT_TIMEOUT,
+        sleep_threshold: float = SLEEP_THRESHOLD
+    ):
         self.is_connected.wait(self.WAIT_TIMEOUT)
 
-        try:
-            return self._send(data, timeout=timeout)
-        except (OSError, TimeoutError, InternalServerError) as e:
-            if retries == 0:
-                raise e from None
+        if isinstance(data, (functions.InvokeWithoutUpdates, functions.InvokeWithTakeout)):
+            query = data.query
+        else:
+            query = data
 
-            (log.warning if retries < 2 else log.info)(
-                "[{}] Retrying {} due to {}".format(
-                    Session.MAX_RETRIES - retries + 1,
-                    data.QUALNAME, e))
+        query = ".".join(query.QUALNAME.split(".")[1:])
 
-            time.sleep(0.5)
-            return self.send(data, retries - 1, timeout)
+        while True:
+            try:
+                return self._send(data, timeout=timeout)
+            except FloodWait as e:
+                amount = e.x
+
+                if amount > sleep_threshold:
+                    raise
+
+                log.warning('[{}] Sleeping for {}s (required by "{}")'.format(
+                    self.client.session_name, amount, query))
+
+                time.sleep(amount)
+            except (OSError, TimeoutError, InternalServerError) as e:
+                if retries == 0:
+                    raise e from None
+
+                (log.warning if retries < 2 else log.info)(
+                    '[{}] Retrying "{}" due to {}'.format(
+                        Session.MAX_RETRIES - retries + 1,
+                        query, e))
+
+                time.sleep(0.5)
+
+                return self.send(data, retries - 1, timeout)