Improve websockets performance by using picows when available

README.rst

@@ -310,6 +310,17 @@ Python-binance also supports `orjson` for parsing JSON since it is much faster t
 
 However, `orjson` is not enabled by default because it is not supported by every python interpreter. If you want to opt-in, you just need to install it (`pip install orjson`) on your local environment. Python-binance will detect the installion and pick it up automatically.
 
+Faster websockets with picows
+-----------------------------
+Python-binance supports `picows` as a faster alternative to `websockets` library.
+
+It is not enabled by default. If you want to opt-in, you just need to install picows (any version starting from 2.1)
+
+.. code:: sh
+
+    $ pip install picows>=2.1
+
+
 LLM & AI Agent Support
 ----------------------
 

binance/ws/reconnecting_websocket.py

@@ -14,10 +14,7 @@
 except ImportError:
     pass
 
-try:
-    from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK  # type: ignore
-except ImportError:
-    from websockets import ConnectionClosedError, ConnectionClosedOK  # type: ignore
+from .websockets_compat import ConnectionClosedError, ConnectionClosedOK  # type: ignore
 
 
 Proxy = None
@@ -30,7 +27,7 @@
 except ImportError:
     pass
 
-import websockets as ws
+from .websockets_compat import websockets as ws
 
 from binance.exceptions import (
     BinanceWebsocketClosed,

binance/ws/websocket_api.py

@@ -1,8 +1,7 @@
 from typing import Dict, Optional
 import asyncio
 
-from websockets import WebSocketClientProtocol  # type: ignore
-
+from .websockets_compat import websockets, WebSocketClientProtocol
 from .constants import WSListenerState
 from .reconnecting_websocket import ReconnectingWebsocket
 from binance.exceptions import BinanceAPIException, BinanceWebsocketUnableToConnect

binance/ws/websockets_compat.py

@@ -0,0 +1,40 @@
+import re
+
+
+def _check_picows_version():
+    import picows
+
+    match = re.match(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?", picows.__version__)
+    if not match:
+        raise ImportError("picows>=2.1.0 is required")
+
+    version = tuple(int(part or 0) for part in match.groups())
+
+    MIN_PICOWS_VERSION = (2, 1, 0)
+    if not version >= MIN_PICOWS_VERSION:
+        raise ImportError("picows>=2.1.0 is required")
+
+
+try:
+    _check_picows_version()
+
+    import picows.websockets as websockets
+    from picows.websockets import (
+        ConnectionClosed,
+        ConnectionClosedError,
+        ConnectionClosedOK,
+        State,
+        WebSocketClientProtocol,
+        protocol as ws_protocol,
+    )
+    websockets_package_name = "picows.websockets"
+except ImportError:
+    import websockets
+    from websockets import protocol as ws_protocol, WebSocketClientProtocol
+    State = ws_protocol.State
+    try:
+        from websockets.exceptions import ConnectionClosed, ConnectionClosedError, ConnectionClosedOK  # type: ignore
+    except ImportError:
+        from websockets import ConnectionClosed, ConnectionClosedError, ConnectionClosedOK  # type: ignore
+
+    websockets_package_name = "websockets"

tests/test_error_propagation.py

@@ -10,9 +10,7 @@
 import pytest
 from unittest.mock import AsyncMock, PropertyMock
 
-from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
-import websockets.protocol as ws_protocol
-
+from binance.ws.websockets_compat import ws_protocol, ConnectionClosedError, ConnectionClosedOK
 from binance.ws.reconnecting_websocket import ReconnectingWebsocket
 from binance.ws.websocket_api import WebsocketAPI
 from binance.ws.constants import WSListenerState

tests/test_reconnecting_websocket.py

@@ -3,11 +3,10 @@
 import gzip
 import json
 from unittest.mock import patch, create_autospec, Mock
+from binance.ws.websockets_compat import WebSocketClientProtocol, State, websockets_package_name  # type: ignore
 from binance.ws.reconnecting_websocket import ReconnectingWebsocket
 from binance.ws.constants import WSListenerState
 from binance.exceptions import BinanceWebsocketUnableToConnect, ReadLoopClosed
-from websockets import WebSocketClientProtocol  # type: ignore
-from websockets.protocol import State
 import asyncio
 
 try:
@@ -132,7 +131,7 @@ async def test_recieve_invalid_json():
     mock_socket.state = AsyncMock()
 
     # Mock websockets.connect to return our mock socket
-    with patch("websockets.connect") as mock_connect:
+    with patch(f"{websockets_package_name}.connect") as mock_connect:
         mock_connect.return_value.__aenter__.return_value = mock_socket
 
         ws = ReconnectingWebsocket(url="wss://test.url")
@@ -152,7 +151,7 @@ async def test_receive_valid_json():
     mock_socket.state = AsyncMock()
 
     # Mock websockets.connect to return our mock socket
-    with patch("websockets.connect") as mock_connect:
+    with patch(f"{websockets_package_name}.connect") as mock_connect:
         mock_connect.return_value.__aenter__.return_value = mock_socket
 
         ws = ReconnectingWebsocket(url="wss://test.url")
@@ -188,7 +187,7 @@ async def test_connect_fails_to_connect_after_disconnect():
         Exception("Connection failed"),  # Subsequent calls fail
     ]
 
-    with patch("websockets.connect", return_value=mock_connect.return_value):
+    with patch(f"{websockets_package_name}.connect", return_value=mock_connect.return_value):
         ws = ReconnectingWebsocket(url="wss://test.url")
         async with ws as ws:
             assert ws.ws is not None

tests/test_threaded_stream.py

@@ -1,7 +1,7 @@
 import pytest
 import asyncio
 
-import websockets
+from binance.ws.websockets_compat import ConnectionClosed
 from binance.ws.threaded_stream import ThreadedApiManager
 from unittest.mock import Mock
 
@@ -71,7 +71,7 @@ async def controlled_recv():
         recv_count += 1
         # If we've stopped the socket or read enough times, simulate connection closing
         if not manager._socket_running.get(socket_name) or recv_count > 2:
-            raise websockets.exceptions.ConnectionClosed(None, None)
+            raise ConnectionClosed(None, None)
         await asyncio.sleep(0.1)
         return '{"e": "value"}'
 
@@ -95,7 +95,7 @@ async def controlled_recv():
     # Wait for the listener task to complete
     try:
         await asyncio.wait_for(listener_task, timeout=1.0)
-    except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
+    except (asyncio.TimeoutError, ConnectionClosed):
         pass  # These exceptions are expected during shutdown
 
     assert socket_name not in manager._socket_running
@@ -134,7 +134,7 @@ async def controlled_recv():
     # Wait for the listener to finish
     try:
         await asyncio.wait_for(listener_task, timeout=1.0)
-    except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
+    except (asyncio.TimeoutError, ConnectionClosed):
         listener_task.cancel()
 
     # Callback should not have been called (no successful messages)