From 69d1432251c15a150d639d2a547729b5d3ee105a Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 24 Jun 2018 17:47:29 +0200
Subject: [PATCH] Allow user-defined values for retries and timeout

---
 pyrogram/client/client.py          | 10 ++++++++--
 pyrogram/client/ext/base_client.py |  3 ++-
 pyrogram/session/session.py        | 10 +++++-----
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index 2312d6dd..1897e8d3 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -817,7 +817,7 @@ class Client(Methods, BaseClient):
 
         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 @@ class Client(Methods, BaseClient):
             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
index 578cf168..d8e30283 100644
--- a/pyrogram/client/ext/base_client.py
+++ b/pyrogram/client/ext/base_client.py
@@ -23,6 +23,7 @@ from threading import Lock
 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 @@ class BaseClient:
 
         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
index 7e90cfff..d8c8651d 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -391,7 +391,7 @@ class Session:
 
         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 @@ class Session:
             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 @@ class Session:
             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 @@ class Session:
                     datetime.now(), type(data)))
 
             time.sleep(0.5)
-            return self.send(data, retries - 1)
+            return self.send(data, retries - 1, timeout)