aiohttp.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. from __future__ import annotations
  2. from aiohttp import ClientSession, ClientResponse, ClientTimeout, BaseConnector, FormData
  3. from typing import AsyncIterator, Any, Optional
  4. from .defaults import DEFAULT_HEADERS
  5. from ..errors import MissingRequirementsError
  6. class StreamResponse(ClientResponse):
  7. async def iter_lines(self) -> AsyncIterator[bytes]:
  8. async for line in self.content:
  9. yield line.rstrip(b"\r\n")
  10. async def iter_content(self) -> AsyncIterator[bytes]:
  11. async for chunk in self.content.iter_any():
  12. yield chunk
  13. async def json(self, content_type: str = None) -> Any:
  14. return await super().json(content_type=content_type)
  15. class StreamSession(ClientSession):
  16. def __init__(
  17. self,
  18. headers: dict = {},
  19. timeout: int = None,
  20. connector: BaseConnector = None,
  21. proxy: str = None,
  22. proxies: dict = {},
  23. impersonate = None,
  24. **kwargs
  25. ):
  26. if impersonate:
  27. headers = {
  28. **DEFAULT_HEADERS,
  29. **headers
  30. }
  31. connect = None
  32. if isinstance(timeout, tuple):
  33. connect, timeout = timeout;
  34. if timeout is not None:
  35. timeout = ClientTimeout(timeout, connect)
  36. if proxy is None:
  37. proxy = proxies.get("all", proxies.get("https"))
  38. super().__init__(
  39. **kwargs,
  40. timeout=timeout,
  41. response_class=StreamResponse,
  42. connector=get_connector(connector, proxy),
  43. headers=headers
  44. )
  45. def get_connector(connector: BaseConnector = None, proxy: str = None, rdns: bool = False) -> Optional[BaseConnector]:
  46. if proxy and not connector:
  47. try:
  48. from aiohttp_socks import ProxyConnector
  49. if proxy.startswith("socks5h://"):
  50. proxy = proxy.replace("socks5h://", "socks5://")
  51. rdns = True
  52. connector = ProxyConnector.from_url(proxy, rdns=rdns)
  53. except ImportError:
  54. raise MissingRequirementsError('Install "aiohttp_socks" package for proxy support')
  55. return connector