mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-22 18:07:21 +00:00
Code Issues Releases Wiki Activity

Update connection.py

Browse Source
This commit is contained in:
Dan 2022-12-26 16:29:05 +01:00
parent dd4e41f63f
commit 7182a7cff7
1 changed files with 15 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
pyrogram/connection/connection.py
View File

@ -20,67 +20,53 @@ import asyncio
import logging import logging
from typing import Optional from typing import Optional
from .transport import * from .transport import TCP, TCPAbridgedO
from ..session.internals import DataCenter from ..session.internals import DataCenter
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
class Connection: class Connection:
MAX_RETRIES = 3 MAX_CONNECTION_ATTEMPTS = 3
MODES = { def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict, media: bool = False):
0: TCPFull,
1: TCPAbridged,
2: TCPIntermediate,
3: TCPAbridgedO,
4: TCPIntermediateO
}
def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict, media: bool = False, mode: int = 3):
self.dc_id = dc_id self.dc_id = dc_id
self.test_mode = test_mode self.test_mode = test_mode
self.ipv6 = ipv6 self.ipv6 = ipv6
self.proxy = proxy self.proxy = proxy
self.media = media self.media = media
self.address = DataCenter(dc_id, test_mode, ipv6, media)
self.mode = self.MODES.get(mode, TCPAbridged)
self.protocol = None # type: TCP self.address = DataCenter(dc_id, test_mode, ipv6, media)
self.protocol: TCP = None
async def connect(self): async def connect(self):
for i in range(Connection.MAX_RETRIES): for i in range(Connection.MAX_CONNECTION_ATTEMPTS):
self.protocol = self.mode(self.ipv6, self.proxy) self.protocol = TCPAbridgedO(self.ipv6, self.proxy)
try: try:
log.info("Connecting...") log.info("Connecting...")
await self.protocol.connect(self.address) await self.protocol.connect(self.address)
except OSError as e: except OSError as e:
log.warning(f"Unable to connect due to network issues: {e}") log.warning("Unable to connect due to network issues: %s", e)
await self.protocol.close() await self.protocol.close()
await asyncio.sleep(1) await asyncio.sleep(1)
else: else:
log.info("Connected! {} DC{}{} - IPv{} - {}".format( log.info("Connected! %s DC%s%s - IPv%s",
"Test" if self.test_mode else "Production", "Test" if self.test_mode else "Production",
self.dc_id, self.dc_id,
" (media)" if self.media else "", " (media)" if self.media else "",
"6" if self.ipv6 else "4", "6" if self.ipv6 else "4")
self.mode.__name__,
))
break break
else: else:
log.warning("Connection failed! Trying again...") log.warning("Connection failed! Trying again...")
raise TimeoutError raise ConnectionError
async def close(self): async def close(self):
await self.protocol.close() await self.protocol.close()
log.info("Disconnected") log.info("Disconnected")
async def send(self, data: bytes): async def send(self, data: bytes):
try: await self.protocol.send(data)
await self.protocol.send(data)
except Exception as e:
raise OSError(e)
async def recv(self) -> Optional[bytes]: async def recv(self) -> Optional[bytes]:
return await self.protocol.recv() return await self.protocol.recv()