From 4f6990d735d63404ff7c5297e9c7b12d2f76599b Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 3 Jan 2019 12:20:42 +0100
Subject: [PATCH] Add takeout parameter in Client This lets the client use a
 takeout session instead of a normal one. Takeout sessions are useful for
 exporting Telegram data. Methods invoked inside a takeout session are less
 prone to throw FloodWait exceptions.

---
 pyrogram/client/client.py          | 22 +++++++++++++++++++++-
 pyrogram/client/ext/base_client.py |  2 ++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index daef9d34..bd72c582 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -157,6 +157,13 @@ class Client(Methods, BaseClient):
             When updates are disabled your client can't receive any new message.
             Useful for batch programs that don't need to deal with updates.
             Defaults to False (updates enabled and always received).
+
+        takeout (``bool``, *optional*):
+            Pass True to let the client use a takeout session instead of a normal one, implies no_updates.
+            Useful for exporting your Telegram data. Methods invoked inside a takeout session (such as get_history,
+            download_media, ...) are less prone to throw FloodWait exceptions.
+            Only available for users, bots will ignore this parameter.
+            Defaults to False (normal session).
     """
 
     def __init__(self,
@@ -180,7 +187,8 @@ class Client(Methods, BaseClient):
                  workdir: str = BaseClient.WORKDIR,
                  config_file: str = BaseClient.CONFIG_FILE,
                  plugins_dir: str = None,
-                 no_updates: bool = None):
+                 no_updates: bool = None,
+                 takeout: bool = None):
         super().__init__()
 
         self.session_name = session_name
@@ -205,6 +213,7 @@ class Client(Methods, BaseClient):
         self.config_file = config_file
         self.plugins_dir = plugins_dir
         self.no_updates = no_updates
+        self.takeout = takeout
 
         self.dispatcher = Dispatcher(self, workers)
 
@@ -261,6 +270,10 @@ class Client(Methods, BaseClient):
                 self.save_session()
 
             if self.bot_token is None:
+                if self.takeout:
+                    self.takeout_id = self.send(functions.account.InitTakeoutSession()).id
+                    log.warning("Takeout session {} initiated".format(self.takeout_id))
+
                 now = time.time()
 
                 if abs(now - self.date) > Client.OFFLINE_SLEEP:
@@ -316,6 +329,10 @@ class Client(Methods, BaseClient):
         if not self.is_started:
             raise ConnectionError("Client is already stopped")
 
+        if self.takeout_id:
+            self.send(functions.account.FinishTakeoutSession())
+            log.warning("Takeout session {} finished".format(self.takeout_id))
+
         Syncer.remove(self)
         self.dispatcher.stop()
 
@@ -954,6 +971,9 @@ class Client(Methods, BaseClient):
         if self.no_updates:
             data = functions.InvokeWithoutUpdates(data)
 
+        if self.takeout_id:
+            data = functions.InvokeWithTakeout(self.takeout_id, data)
+
         r = self.session.send(data, retries, timeout)
 
         self.fetch_peers(getattr(r, "users", []))
diff --git a/pyrogram/client/ext/base_client.py b/pyrogram/client/ext/base_client.py
index 7aee5302..cecbf5fb 100644
--- a/pyrogram/client/ext/base_client.py
+++ b/pyrogram/client/ext/base_client.py
@@ -90,6 +90,8 @@ class BaseClient:
         self.is_started = None
         self.is_idle = None
 
+        self.takeout_id = None
+
         self.updates_queue = Queue()
         self.updates_workers_list = []
         self.download_queue = Queue()