summaryrefslogtreecommitdiffstats
path: root/g4f/requests/aiohttp.py
blob: d9bd6541a8eb98bde6f8bf433d9f53185d9e4258 (plain) (blame)
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
from __future__ import annotations

from aiohttp import ClientSession, ClientResponse, ClientTimeout
from typing import AsyncGenerator, Any

from ..providers.helper import get_connector
from .defaults import DEFAULT_HEADERS

class StreamResponse(ClientResponse):
    async def iter_lines(self) -> AsyncGenerator[bytes, None]:
        async for line in self.content:
            yield line.rstrip(b"\r\n")

    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
        )