1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
from __future__ import annotations
from aiohttp import ClientSession, ClientResponse, ClientTimeout, BaseConnector, FormData
from typing import AsyncIterator, Any, Optional
from .defaults import DEFAULT_HEADERS
from ..errors import MissingRequirementsError
class StreamResponse(ClientResponse):
async def iter_lines(self) -> AsyncIterator[bytes]:
async for line in self.content:
yield line.rstrip(b"\r\n")
async def iter_content(self) -> AsyncIterator[bytes]:
async for chunk in self.content.iter_any():
yield chunk
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 = {
**DEFAULT_HEADERS,
**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
)
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
|
