From 628b7391e38b837842b5f7a4075e064753f4e8f6 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 20 Mar 2017 12:23:53 +0100
Subject: [PATCH] Greatly improve the handling of working with different DC

---
 telethon/network/mtproto_sender.py |  2 +-
 telethon/telegram_client.py        | 33 ++++++++++++++++--------------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/telethon/network/mtproto_sender.py b/telethon/network/mtproto_sender.py
index eddaa8fe..d9624a3f 100755
--- a/telethon/network/mtproto_sender.py
+++ b/telethon/network/mtproto_sender.py
@@ -328,7 +328,7 @@ class MtProtoSender:
                     error.additional_data))
                 sleep(error.additional_data)
 
-            elif error.message.startswith('PHONE_MIGRATE_'):
+            elif '_MIGRATE_' in error.message:
                 raise InvalidDCError(error.additional_data)
 
             else:
diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 510c8a8c..bb23284b 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -152,17 +152,28 @@ class TelegramClient:
 
     # region Telegram requests functions
 
-    def invoke(self, request, timeout=timedelta(seconds=5)):
+    def invoke(self, request, timeout=timedelta(seconds=5), throw_invalid_dc=False):
         """Invokes a MTProtoRequest (sends and receives it) and returns its result.
            An optional timeout can be given to cancel the operation after the time delta.
-           Timeout can be set to None for no timeout"""
+           Timeout can be set to None for no timeout.
+
+           If throw_invalid_dc is True, these errors won't be caught (useful to
+           avoid infinite recursion). This should not be set to True manually."""
         if not issubclass(type(request), MTProtoRequest):
             raise ValueError('You can only invoke MtProtoRequests')
 
-        self.sender.send(request)
-        self.sender.receive(request, timeout)
+        try:
+            self.sender.send(request)
+            self.sender.receive(request, timeout)
 
-        return request.result
+            return request.result
+
+        except InvalidDCError as error:
+            if throw_invalid_dc:
+                raise error
+
+            self.reconnect_to_dc(error.new_dc)
+            return self.invoke(request, timeout=timeout, throw_invalid_dc=True)
 
     # region Authorization requests
 
@@ -173,16 +184,8 @@ class TelegramClient:
 
     def send_code_request(self, phone_number):
         """Sends a code request to the specified phone number"""
-        request = SendCodeRequest(phone_number, self.api_id, self.api_hash)
-        completed = False
-        while not completed:
-            try:
-                result = self.invoke(request)
-                self.phone_code_hashes[phone_number] = result.phone_code_hash
-                completed = True
-
-            except InvalidDCError as error:
-                self.reconnect_to_dc(error.new_dc)
+        result = self.invoke(SendCodeRequest(phone_number, self.api_id, self.api_hash))
+        self.phone_code_hashes[phone_number] = result.phone_code_hash
 
     def sign_in(self, phone_number=None, code=None, password=None):
         """Completes the authorization of a phone number by providing the received code.