chore: remove redundant of proxy_url schema for socks version

This commit is contained in:
mpa 2020-03-19 06:46:19 +04:00
parent c1adbb02b9
commit 1f8fa0c4e8
2 changed files with 23 additions and 21 deletions

View file

@ -1,6 +1,6 @@
from __future__ import annotations from __future__ import annotations
from typing import AsyncGenerator, Callable, Optional, TypeVar, Tuple, Dict, Any, cast from typing import AsyncGenerator, Callable, Optional, TypeVar, Tuple, Dict, Any, Union, cast
from aiohttp import ClientSession, ClientTimeout, FormData, BasicAuth, TCPConnector from aiohttp import ClientSession, ClientTimeout, FormData, BasicAuth, TCPConnector
@ -9,16 +9,16 @@ from aiogram.api.methods import Request, TelegramMethod
from .base import PRODUCTION, BaseSession, TelegramAPIServer from .base import PRODUCTION, BaseSession, TelegramAPIServer
T = TypeVar("T") T = TypeVar("T")
_ProxyType = Tuple[str, BasicAuth] _ProxyType = Union[str, Tuple[str, BasicAuth]]
class AiohttpSession(BaseSession[_ProxyType]): class AiohttpSession(BaseSession[_ProxyType]):
def __init__( def __init__(
self, self,
api: TelegramAPIServer = PRODUCTION, api: TelegramAPIServer = PRODUCTION,
proxy: Optional[_ProxyType] = None,
json_loads: Optional[Callable] = None, json_loads: Optional[Callable] = None,
json_dumps: Optional[Callable] = None, json_dumps: Optional[Callable] = None,
proxy: Optional[_ProxyType] = None,
): ):
super(AiohttpSession, self).__init__( super(AiohttpSession, self).__init__(
api=api, api=api,
@ -31,34 +31,36 @@ class AiohttpSession(BaseSession[_ProxyType]):
self.cfg.connector_init = cast(Dict[str, Any], {}) self.cfg.connector_init = cast(Dict[str, Any], {})
if self.proxy: if self.proxy:
proxy_url, proxy_auth = self.proxy
try: try:
from aiohttp_socks import ProxyConnector from aiohttp_socks import ProxyConnector
from aiohttp_socks.utils import parse_proxy_url from aiohttp_socks.utils import parse_proxy_url
except ImportError as exc: except ImportError as exc:
raise UserWarning( raise UserWarning(
"In order to use aiohttp client for proxy requests, install " "In order to use aiohttp client for proxy requests, install "
"https://pypi.org/project/aiohttp-socks/0.3.4" "https://pypi.org/project/aiohttp-socks/"
) from exc ) from exc
if proxy_url.startswith('socks5://') or proxy_url.startswith('socks4://'): if isinstance(self.proxy, str):
self.cfg.connector_type = ProxyConnector proxy_url, proxy_auth = self.proxy, None
else:
proxy_url, proxy_auth = self.proxy
proxy_type, host, port, username, password = parse_proxy_url(proxy_url) self.cfg.connector_type = ProxyConnector
if proxy_auth:
if not username:
username = proxy_auth.login
if not password:
password = proxy_auth.password
self.cfg.connector_init.update( proxy_type, host, port, username, password = parse_proxy_url(proxy_url)
dict( if proxy_auth:
proxy_type=proxy_type, host=host, port=port, if not username:
username=username, password=password, username = proxy_auth.login
rdns=True, if not password:
) password = proxy_auth.password
self.cfg.connector_init.update(
dict(
proxy_type=proxy_type, host=host, port=port,
username=username, password=password,
rdns=True,
) )
)
async def create_session(self) -> ClientSession: async def create_session(self) -> ClientSession:
if self._session is None or self._session.closed: if self._session is None or self._session.closed:

View file

@ -19,9 +19,9 @@ class BaseSession(abc.ABC, Generic[_ProxyType]):
def __init__( def __init__(
self, self,
api: Optional[TelegramAPIServer] = None, api: Optional[TelegramAPIServer] = None,
proxy: Optional[_ProxyType] = None,
json_loads: Optional[Callable[[Any], Any]] = None, json_loads: Optional[Callable[[Any], Any]] = None,
json_dumps: Optional[Callable[[Any], Any]] = None, json_dumps: Optional[Callable[[Any], Any]] = None,
proxy: Optional[_ProxyType] = None,
) -> None: ) -> None:
if api is None: if api is None:
api = PRODUCTION api = PRODUCTION