refactor: Split WebSocket support into separate sync/async implementations

README.md

@@ -200,28 +200,38 @@ async with AsyncSession() as s:
 ### WebSockets
 
 ```python
-from curl_cffi.requests import Session, WebSocket
+from curl_cffi.requests import WebSocket
 
-def on_message(ws: WebSocket, message):
+def on_message(ws: WebSocket, message: str | bytes):
     print(message)
 
-with Session() as s:
-    ws = s.ws_connect(
-        "wss://api.gemini.com/v1/marketdata/BTCUSD",
-        on_message=on_message,
-    )
-    ws.run_forever()
+ws = WebSocket(on_message=on_message)
+ws.run_forever("wss://api.gemini.com/v1/marketdata/BTCUSD")
 ```
 
 For low-level APIs, Scrapy integration and other advanced topics, see the
 [docs](https://curl-cffi.readthedocs.io) for more details.
 
+### asyncio WebSockets
+
+```python
+import asyncio
+from curl_cffi.requests import AsyncSession
+
+async with AsyncSession() as s:
+    ws = await s.ws_connect("wss://echo.websocket.org")
+    await asyncio.gather(*[ws.send_str("Hello, World!") for _ in range(10)])
+    async for message in ws:
+        print(message)
+```
+
 ## Acknowledgement
 
 - Originally forked from [multippt/python_curl_cffi](https://github.com/multippt/python_curl_cffi), which is under the MIT license.
 - Headers/Cookies files are copied from [httpx](https://github.com/encode/httpx/blob/master/httpx/_models.py), which is under the BSD license.
 - Asyncio support is inspired by Tornado's curl http client.
-- The WebSocket API is inspired by [websocket_client](https://github.com/websocket-client/websocket-client).
+- The synchronous WebSocket API is inspired by [websocket_client](https://github.com/websocket-client/websocket-client).
+- The asynchronous WebSocket API is inspired by [aiohttp](https://github.com/aio-libs/aiohttp).
 
 
 ## Sponsor

curl_cffi/curl.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import re
 import warnings
 from http.cookies import SimpleCookie
 from pathlib import Path
-from typing import Any, List, Literal, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast
 
 import certifi
 
@@ -13,6 +15,15 @@
 REASON_PHRASE_RE = re.compile(rb"HTTP/\d\.\d [0-9]{3} (.*)")
 STATUS_LINE_RE = re.compile(rb"HTTP/(\d\.\d) ([0-9]{3}) (.*)")
 
+if TYPE_CHECKING:
+
+    class CurlWsFrame:
+        age: int
+        flags: int
+        offset: int
+        bytesleft: int
+        len: int
+
 
 class CurlError(Exception):
     """Base exception for curl_cffi package"""
@@ -234,7 +245,7 @@ def getinfo(self, option: CurlInfo) -> Union[bytes, int, float, List]:
             0x200000: "long*",
             0x300000: "double*",
             0x400000: "struct curl_slist **",
-            0x500000: "long*"
+            0x500000: "long*",
         }
         ret_cast_option = {
             0x100000: ffi.string,
@@ -310,7 +321,7 @@ def clean_after_perform(self, clear_headers: bool = True) -> None:
                 lib.curl_slist_free_all(self._proxy_headers)
             self._proxy_headers = ffi.NULL
 
-    def duphandle(self) -> "Curl":
+    def duphandle(self) -> Curl:
         """Wrapper for ``curl_easy_duphandle``.
 
         This is not a full copy of entire curl object in python. For example, headers
@@ -379,7 +390,7 @@ def close(self) -> None:
         ffi.release(self._error_buffer)
         self._resolve = ffi.NULL
 
-    def ws_recv(self, n: int = 1024) -> Tuple[bytes, Any]:
+    def ws_recv(self, n: int = 1024) -> Tuple[bytes, CurlWsFrame]:
         """Receive a frame from a websocket connection.
 
         Args:
@@ -422,10 +433,6 @@ def ws_send(self, payload: bytes, flags: CurlWsFlag = CurlWsFlag.BINARY) -> int:
         self._check_error(ret, "WS_SEND")
         return n_sent[0]
 
-    def ws_close(self) -> None:
-        """Send the close frame."""
-        self.ws_send(b"", CurlWsFlag.CLOSE)
-
 
 class CurlMime:
     """Wrapper for the ``curl_mime_`` API."""

curl_cffi/requests/__init__.py

@@ -17,6 +17,7 @@
     "Headers",
     "Request",
     "Response",
+    "AsyncWebSocket",
     "WebSocket",
     "WebSocketError",
     "WsCloseCode",
@@ -38,7 +39,7 @@
 from .impersonate import BrowserType, BrowserTypeLiteral, ExtraFingerprints, ExtraFpDict
 from .models import Request, Response
 from .session import AsyncSession, HttpMethod, ProxySpec, Session, ThreadType
-from .websockets import WebSocket, WebSocketError, WsCloseCode
+from .websockets import AsyncWebSocket, WebSocket, WebSocketError, WsCloseCode
 
 
 def request(

curl_cffi/requests/cookies.py

@@ -273,7 +273,9 @@ def get_dict(self, domain: Optional[str] = None, path: Optional[str] = None) ->
         """
         ret = {}
         for cookie in self.jar:
-            if (domain is None or cookie.domain == domain) and (path is None or cookie.path == path):
+            if (domain is None or cookie.domain == domain) and (
+                path is None or cookie.path == path
+            ):
                 ret[cookie.name] = cookie.value
         return ret