diff --git a/README.rst b/README.rst index 48efec45..be58302f 100755 --- a/README.rst +++ b/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 ---------------------- diff --git a/binance/ws/reconnecting_websocket.py b/binance/ws/reconnecting_websocket.py index 9f25af70..d93916eb 100644 --- a/binance/ws/reconnecting_websocket.py +++ b/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, diff --git a/binance/ws/websocket_api.py b/binance/ws/websocket_api.py index ae62e788..09fa63d4 100644 --- a/binance/ws/websocket_api.py +++ b/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 diff --git a/binance/ws/websockets_compat.py b/binance/ws/websockets_compat.py new file mode 100644 index 00000000..2fc509fa --- /dev/null +++ b/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" diff --git a/tests/test_error_propagation.py b/tests/test_error_propagation.py index d419c5ed..c2273b39 100644 --- a/tests/test_error_propagation.py +++ b/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 diff --git a/tests/test_reconnecting_websocket.py b/tests/test_reconnecting_websocket.py index 7433a411..51820b41 100644 --- a/tests/test_reconnecting_websocket.py +++ b/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 diff --git a/tests/test_threaded_stream.py b/tests/test_threaded_stream.py index 44a62a52..dfb539af 100644 --- a/tests/test_threaded_stream.py +++ b/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)