From a427465231c51417d31fbc1b46902ae92d335bd3 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Tue, 30 May 2017 12:14:29 +0200
Subject: [PATCH] Tidy up working with different DCs

---
 telethon/telegram_client.py | 42 +++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 16 deletions(-)

diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 80efc6c3..21b1ba03 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -146,22 +146,6 @@ class TelegramClient:
                                  .format(error))
             return False
 
-    def _reconnect_to_dc(self, dc_id):
-        """Reconnects to the specified DC ID. This is automatically called after an InvalidDCError is raised"""
-        if self.dc_options is None or not self.dc_options:
-            raise ConnectionError(
-                "Can't reconnect. Stabilise an initial connection first.")
-
-        dc = next(dc for dc in self.dc_options if dc.id == dc_id)
-
-        self.transport.close()
-        self.transport = None
-        self.session.server_address = dc.ip_address
-        self.session.port = dc.port
-        self.session.save()
-
-        self.connect(reconnect=True)
-
     def disconnect(self):
         """Disconnects from the Telegram server and stops all the spawned threads"""
         self._set_updates_thread(running=False)
@@ -179,6 +163,32 @@ class TelegramClient:
 
     # endregion
 
+    # region Working with different Data Centers
+
+    def _reconnect_to_dc(self, dc_id):
+        """Reconnects to the specified DC ID. This is automatically
+           called after an InvalidDCError is raised"""
+        dc = self._get_dc(dc_id)
+
+        self.transport.close()
+        self.transport = None
+        self.session.server_address = dc.ip_address
+        self.session.port = dc.port
+        self.session.save()
+
+        self.connect(reconnect=True)
+
+    def _get_dc(self, dc_id):
+        """Gets the Data Center (DC) associated to 'dc_id'"""
+        if not self.dc_options:
+            raise ConnectionError(
+                'Cannot determine the required data center IP address. '
+                'Stabilise an initial connection first.')
+
+        return next(dc for dc in self.dc_options if dc.id == dc_id)
+
+    # endregion
+
     # region Telegram requests functions
 
     def invoke(self, request, timeout=timedelta(seconds=5), throw_invalid_dc=False):