From fdc8b88d17f7b9fa59aafa02612017c8cb32dbf9 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Sat, 20 Apr 2024 20:06:35 +0200 Subject: Improve cli argparser, Improve curl_cffi websocket support --- g4f/requests/curl_cffi.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) (limited to 'g4f/requests') diff --git a/g4f/requests/curl_cffi.py b/g4f/requests/curl_cffi.py index d0d44ba7..1464cb32 100644 --- a/g4f/requests/curl_cffi.py +++ b/g4f/requests/curl_cffi.py @@ -79,10 +79,10 @@ class StreamSession(AsyncSession): return StreamResponse(super().request(method, url, stream=True, **kwargs)) def ws_connect(self, url, *args, **kwargs): - return WebSocket(self, url) + return WebSocket(self, url, **kwargs) - def _ws_connect(self, url): - return super().ws_connect(url) + def _ws_connect(self, url, **kwargs): + return super().ws_connect(url, **kwargs) # Defining HTTP methods as partial methods of the request method. head = partialmethod(request, "HEAD") @@ -102,20 +102,22 @@ else: raise RuntimeError("CurlMimi in curl_cffi is missing | pip install -U g4f[curl_cffi]") class WebSocket(): - def __init__(self, session, url) -> None: + def __init__(self, session, url, **kwargs) -> None: if not has_curl_ws: raise RuntimeError("CurlWsFlag in curl_cffi is missing | pip install -U g4f[curl_cffi]") self.session: StreamSession = session self.url: str = url + del kwargs["autoping"] + self.options: dict = kwargs async def __aenter__(self): - self.inner = await self.session._ws_connect(self.url) + self.inner = await self.session._ws_connect(self.url, **self.options) return self async def __aexit__(self, *args): - self.inner.aclose() + await self.inner.aclose() - async def receive_str(self) -> str: + async def receive_str(self, **kwargs) -> str: bytes, _ = await self.inner.arecv() return bytes.decode(errors="ignore") -- cgit v1.2.3