Allow user-defined values for retries and timeout

delivrance · delivrance · commit 69d1432251c1 · 2018-06-24T17:47:29.000+02:00
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
@@ -817,7 +817,7 @@ def updates_worker(self):
 
         log.debug("{} stopped".format(name))
 
-    def send(self, data: Object):
+    def send(self, data: Object, retries: int = Session.MAX_RETRIES, timeout: float = Session.WAIT_TIMEOUT):
         """Use this method to send Raw Function queries.
 
         This method makes possible to manually call every single Telegram API method in a low-level manner.
@@ -828,13 +828,19 @@ def send(self, data: Object):
             data (``Object``):
                 The API Scheme function filled with proper arguments.
 
+            retries (``int``):
+                Number of retries.
+
+            timeout (``float``):
+                Timeout in seconds.
+
         Raises:
             :class:`Error <pyrogram.Error>`
         """
         if not self.is_started:
             raise ConnectionError("Client has not been started")
 
-        r = self.session.send(data)
+        r = self.session.send(data, retries, timeout)
 
         self.fetch_peers(getattr(r, "users", []))
         self.fetch_peers(getattr(r, "chats", []))
diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py
@@ -23,6 +23,7 @@
 from ..style import Markdown, HTML
 from ...api.core import Object
 from ...session.internals import MsgId
+from ...session import Session
 
 
 class BaseClient:
@@ -77,7 +78,7 @@ def __init__(self):
 
         self.disconnect_handler = None
 
-    def send(self, data: Object):
+    def send(self, data: Object, retries: int = Session.MAX_RETRIES, timeout: float = Session.WAIT_TIMEOUT):
         pass
 
     def resolve_peer(self, peer_id: int or str):
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
@@ -391,7 +391,7 @@ def recv(self):
 
         log.debug("RecvThread stopped")
 
-    def _send(self, data: Object, wait_response: bool = True):
+    def _send(self, data: Object, wait_response: bool = True, timeout: float = WAIT_TIMEOUT):
         message = self.msg_factory(data)
         msg_id = message.msg_id
 
@@ -407,7 +407,7 @@ def _send(self, data: Object, wait_response: bool = True):
             raise e
 
         if wait_response:
-            self.results[msg_id].event.wait(self.WAIT_TIMEOUT)
+            self.results[msg_id].event.wait(timeout)
             result = self.results.pop(msg_id).value
 
             if result is None:
@@ -422,11 +422,11 @@ def _send(self, data: Object, wait_response: bool = True):
             else:
                 return result
 
-    def send(self, data: Object, retries: int = MAX_RETRIES):
+    def send(self, data: Object, retries: int = MAX_RETRIES, timeout: float = WAIT_TIMEOUT):
         self.is_connected.wait(self.WAIT_TIMEOUT)
 
         try:
-            return self._send(data)
+            return self._send(data, timeout=timeout)
         except (OSError, TimeoutError, InternalServerError) as e:
             if retries == 0:
                 raise e from None
@@ -437,4 +437,4 @@ def send(self, data: Object, retries: int = MAX_RETRIES):
                     datetime.now(), type(data)))
 
             time.sleep(0.5)
-            return self.send(data, retries - 1)
+            return self.send(data, retries - 1, timeout)
