Implement non-blocking TCP connection (

RknDeveloper · web-flow · commit d2f3cd998a90 · 2024-06-26T10:02:48.000+05:45
diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py
@@ -18,9 +18,9 @@
 
 import asyncio
 import logging
-from typing import Optional
+from typing import Optional, Type
 
-from .transport import TCP, TCPAbridgedO
+from .transport import TCP, TCPAbridged
 from ..session.internals import DataCenter
 
 log = logging.getLogger(__name__)
@@ -29,26 +29,37 @@
 class Connection:
     MAX_CONNECTION_ATTEMPTS = 3
 
-    def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict, media: bool = False):
+    def __init__(
+        self,
+        dc_id: int,
+        test_mode: bool,
+        ipv6: bool,
+        alt_port: bool,
+        proxy: dict,
+        media: bool = False,
+        protocol_factory: Type[TCP] = TCPAbridged
+    ) -> None:
         self.dc_id = dc_id
         self.test_mode = test_mode
         self.ipv6 = ipv6
+        self.alt_port = alt_port
         self.proxy = proxy
         self.media = media
+        self.protocol_factory = protocol_factory
 
-        self.address = DataCenter(dc_id, test_mode, ipv6, media)
-        self.protocol: TCP = None
+        self.address = DataCenter(dc_id, test_mode, ipv6, alt_port, media)
+        self.protocol: Optional[TCP] = None
 
-    async def connect(self):
+    async def connect(self) -> None:
         for i in range(Connection.MAX_CONNECTION_ATTEMPTS):
-            self.protocol = TCPAbridgedO(self.ipv6, self.proxy)
+            self.protocol = self.protocol_factory(ipv6=self.ipv6, proxy=self.proxy)
 
             try:
                 log.info("Connecting...")
                 await self.protocol.connect(self.address)
             except OSError as e:
                 log.warning("Unable to connect due to network issues: %s", e)
-                self.protocol.close()
+                await self.protocol.close()
                 await asyncio.sleep(1)
             else:
                 log.info("Connected! %s DC%s%s - IPv%s",
@@ -59,17 +70,14 @@ async def connect(self):
                 break
         else:
             log.warning("Connection failed! Trying again...")
-            raise TimeoutError
+            raise ConnectionError
 
-    def close(self):
-        self.protocol.close()
+    async def close(self) -> None:
+        await self.protocol.close()
         log.info("Disconnected")
 
-    async def send(self, data: bytes):
-        try:
-            await self.protocol.send(data)
-        except Exception as e:
-            raise OSError(e)
+    async def send(self, data: bytes) -> None:
+        await self.protocol.send(data)
 
     async def recv(self) -> Optional[bytes]:
         return await self.protocol.recv()
