2024-01-26 20:00:44 -05:00
|
|
|
from __future__ import annotations
|
|
|
|
|
2024-03-11 21:06:06 -04:00
|
|
|
from aiohttp import ClientSession, ClientResponse, ClientTimeout, BaseConnector
|
|
|
|
from typing import AsyncIterator, Any, Optional
|
2024-01-26 20:00:44 -05:00
|
|
|
|
2024-01-29 12:14:46 -05:00
|
|
|
from .defaults import DEFAULT_HEADERS
|
2024-03-11 21:06:06 -04:00
|
|
|
from ..errors import MissingRequirementsError
|
2024-01-26 20:00:44 -05:00
|
|
|
|
|
|
|
class StreamResponse(ClientResponse):
|
2024-03-11 21:06:06 -04:00
|
|
|
async def iter_lines(self) -> AsyncIterator[bytes]:
|
2024-01-26 20:00:44 -05:00
|
|
|
async for line in self.content:
|
|
|
|
yield line.rstrip(b"\r\n")
|
|
|
|
|
2024-03-11 21:06:06 -04:00
|
|
|
async def iter_content(self) -> AsyncIterator[bytes]:
|
|
|
|
async for chunk in self.content.iter_any():
|
|
|
|
yield chunk
|
|
|
|
|
2024-01-26 20:00:44 -05:00
|
|
|
async def json(self) -> Any:
|
|
|
|
return await super().json(content_type=None)
|
|
|
|
|
|
|
|
class StreamSession(ClientSession):
|
|
|
|
def __init__(self, headers: dict = {}, timeout: int = None, proxies: dict = {}, impersonate = None, **kwargs):
|
|
|
|
if impersonate:
|
|
|
|
headers = {
|
2024-01-29 12:14:46 -05:00
|
|
|
**DEFAULT_HEADERS,
|
2024-01-26 20:00:44 -05:00
|
|
|
**headers
|
|
|
|
}
|
|
|
|
super().__init__(
|
|
|
|
**kwargs,
|
|
|
|
timeout=ClientTimeout(timeout) if timeout else None,
|
|
|
|
response_class=StreamResponse,
|
|
|
|
connector=get_connector(kwargs.get("connector"), proxies.get("https")),
|
|
|
|
headers=headers
|
2024-03-11 21:06:06 -04:00
|
|
|
)
|
|
|
|
|
|
|
|
def get_connector(connector: BaseConnector = None, proxy: str = None, rdns: bool = False) -> Optional[BaseConnector]:
|
|
|
|
if proxy and not connector:
|
|
|
|
try:
|
|
|
|
from aiohttp_socks import ProxyConnector
|
|
|
|
if proxy.startswith("socks5h://"):
|
|
|
|
proxy = proxy.replace("socks5h://", "socks5://")
|
|
|
|
rdns = True
|
|
|
|
connector = ProxyConnector.from_url(proxy, rdns=rdns)
|
|
|
|
except ImportError:
|
|
|
|
raise MissingRequirementsError('Install "aiohttp_socks" package for proxy support')
|
|
|
|
return connector
|