Use get_session instead of manually creating session

Author: KurimuzonAkuma

pyrogram/client.py

@@ -1355,7 +1355,14 @@ async def get_session(
         if not temporary and sessions.get(dc_id):
             return sessions[dc_id]
 
-        dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6, is_cdn=is_cdn)
+        if not server_address or not port:
+            dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6, is_cdn=is_cdn)
+
+            server_address = server_address or dc_option.ip_address
+            port = port or dc_option.port
+        else:
+            server_address = server_address or await self.storage.server_address()
+            port = port or await self.storage.port()
 
         if is_media:
             auth_key = (await self.get_session(dc_id)).auth_key
@@ -1364,8 +1371,8 @@ async def get_session(
                 auth_key = await Auth(
                     self,
                     dc_id,
-                    server_address or dc_option.ip_address,
-                    port or dc_option.port,
+                    server_address,
+                    port,
                     await self.storage.test_mode()
                 ).create()
             else:
@@ -1374,8 +1381,8 @@ async def get_session(
         session = Session(
             self,
             dc_id,
-            server_address or dc_option.ip_address,
-            port or dc_option.port,
+            server_address,
+            port,
             auth_key,
             await self.storage.test_mode(),
             is_media=is_media

pyrogram/methods/auth/connect.py

@@ -17,7 +17,6 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import pyrogram
-from pyrogram.session import Session
 
 
 class Connect:
@@ -39,17 +38,12 @@ async def connect(
 
         await self.load_session()
 
-        self.session = Session(
-            self,
-            await self.storage.dc_id(),
-            await self.storage.server_address(),
-            await self.storage.port(),
-            await self.storage.auth_key(),
-            await self.storage.test_mode()
+        self.session = self.get_session(
+            server_address=await self.storage.server_address(),
+            port=await self.storage.port(),
+            export_authorization=False,
+            temporary=True
         )
-
-        await self.session.start()
-
         self.is_connected = True
 
         is_ipv6_session = ":" in await self.storage.server_address()

pyrogram/methods/auth/send_code.py

@@ -23,8 +23,6 @@
 from pyrogram import raw
 from pyrogram import types
 from pyrogram.errors import PhoneMigrate, NetworkMigrate
-from pyrogram.raw.base import dc_option
-from pyrogram.session import Session, Auth
 
 log = logging.getLogger(__name__)
 
@@ -105,30 +103,12 @@ async def send_code(
             except (PhoneMigrate, NetworkMigrate) as e:
                 dc_option = await self.get_dc_option(e.value, ipv6=self.ipv6)
                 await self.session.stop()
+                self.session = None
 
                 await self.storage.dc_id(e.value)
                 await self.storage.server_address(dc_option.ip_address)
                 await self.storage.port(dc_option.port)
 
-                await self.storage.auth_key(
-                    await Auth(
-                        self,
-                        await self.storage.dc_id(),
-                        await self.storage.server_address(),
-                        await self.storage.port(),
-                        await self.storage.test_mode()
-                    ).create()
-                )
-
-                self.session = Session(
-                    self,
-                    await self.storage.dc_id(),
-                    await self.storage.server_address(),
-                    await self.storage.port(),
-                    await self.storage.auth_key(),
-                    await self.storage.test_mode()
-                )
-
-                await self.session.start()
+                self.session = await self.get_session(temporary=True, export_authorization=False)
             else:
                 return types.SentCode._parse(r)

pyrogram/methods/auth/sign_in_bot.py

@@ -22,8 +22,6 @@
 from pyrogram import raw
 from pyrogram import types
 from pyrogram.errors import UserMigrate
-from pyrogram.raw.base import dc_option
-from pyrogram.session import Session, Auth
 
 log = logging.getLogger(__name__)
 
@@ -60,29 +58,13 @@ async def sign_in_bot(
             except UserMigrate as e:
                 dc_option = await self.get_dc_option(e.value, ipv6=self.ipv6)
                 await self.session.stop()
+                self.session = None
 
                 await self.storage.dc_id(e.value)
                 await self.storage.server_address(dc_option.ip_address)
                 await self.storage.port(dc_option.port)
-                await self.storage.auth_key(
-                    await Auth(
-                        self,
-                        await self.storage.dc_id(),
-                        await self.storage.server_address(),
-                        await self.storage.port(),
-                        await self.storage.test_mode()
-                    ).create()
-                )
-                self.session = Session(
-                    self,
-                    await self.storage.dc_id(),
-                    await self.storage.server_address(),
-                    await self.storage.port(),
-                    await self.storage.auth_key(),
-                    await self.storage.test_mode()
-                )
 
-                await self.session.start()
+                self.session = self.get_session(temporary=True, export_authorization=False)
             else:
                 await self.storage.user_id(r.user.id)
                 await self.storage.is_bot(True)