import asyncio
import json
from typing import Any, Optional
try:
import websockets
WebSocketClientProtocol = websockets.WebSocketClientProtocol
except ImportError:
WEBSOCKETS_AVAILABLE = False
class WebSocketClientProtocol:
def _raise_import_error(*args, **kwargs):
raise ImportError("websockets module is not installed.")
WebSocketClientProtocol.send = _raise_import_error
WebSocketClientProtocol.close = _raise_import_error
else:
WEBSOCKETS_AVAILABLE = True
class WSProtocol:
closed: bool = True
client: Any = None ws: Optional[WebSocketClientProtocol] = None
url: str
loop: asyncio.AbstractEventLoop
connection_lost_sent: bool
messange_handler: Optional[asyncio.Task] = None
def __init__(
self, url: str, client, loop: asyncio.AbstractEventLoop, ssl_context=None
):
if not WEBSOCKETS_AVAILABLE:
raise ImportError("websockets module is not installed.")
self.client = client
self.url = url
self.loop = loop
self.connection_lost_sent = False
self.ssl_context = ssl_context
async def connect(self):
print(f"Connecting to {self.url}")
try:
if self.ssl_context:
self.ws = await websockets.connect(
self.url, max_size=10 * 1024 * 1024, ssl=self.ssl_context
)
else:
self.ws = await websockets.connect(self.url, max_size=10 * 1024 * 1024)
self.closed = False
self.messange_handler = self.loop.create_task(self.handle_messages())
except Exception as e:
print(f"WebSocket Error: {e}")
self.closed = True
raise e
async def handle_messages(self):
try:
async for msg in self.ws:
print("RECEIVED", msg)
messages = [m for m in msg.split("\n") if m]
for message in messages:
try:
decoded_msg = json.loads(message)
self.client._got_response(decoded_msg)
except ValueError as exc:
print("Bad JSON received from server:", message, exc)
raise
except Exception as e:
print(f"WebSocket Error: {e}")
self.closed = True
self.connection_lost_sent = True
if self.client:
self.client.connection_lost(e)
def send_data(self, data):
msg = json.dumps(data)
print("SENDING", msg)
async def send():
try:
await self.ws.send(msg + "\n")
except Exception as e:
print(f"websocket: closing due to failure to send {msg}: {e}")
self.close()
self.loop.create_task(send())
def close(self):
if not self.closed:
asyncio.create_task(self.ws.close())
self.closed = True
if self.messange_handler:
self.messange_handler.cancel()
self.messange_handler = None
if not self.connection_lost_sent:
self.connection_lost_sent = True
if self.client:
self.client.connection_lost(None)