fa273bcc26
- Add Python toki_socket package (base_client, communicator, tcp/ws client/server) - Add protobuf message definitions for Python - Add go_python.go cross-test implementation and python_go_client - Add agent-task/python_impl plan and code review docs - Update .gitignore for Python cache files with recursive patterns
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import inspect
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from google.protobuf.message import Message
|
|
|
|
try:
|
|
from websockets.asyncio.server import serve as ws_serve
|
|
|
|
_NEW_WEBSOCKETS_API = True
|
|
except ImportError: # websockets 12.x legacy API
|
|
from websockets import serve as ws_serve
|
|
|
|
_NEW_WEBSOCKETS_API = False
|
|
|
|
from toki_socket.communicator import ParserMap
|
|
from toki_socket.ws_client import WsClient
|
|
|
|
|
|
class WsServer:
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
new_client: Callable[[Any], WsClient] | None = None,
|
|
interval_sec: int = 0,
|
|
wait_sec: int = 0,
|
|
parser_map: ParserMap | None = None,
|
|
) -> None:
|
|
self._host = host
|
|
self._port = port
|
|
self._path = path
|
|
self._new_client = new_client
|
|
self._interval_sec = interval_sec
|
|
self._wait_sec = wait_sec
|
|
self._parser_map = parser_map
|
|
self._clients: list[WsClient] = []
|
|
self._server: Any | None = None
|
|
self.on_client_connected: Callable[[WsClient], object] = lambda _: None
|
|
|
|
@property
|
|
def port(self) -> int:
|
|
if self._server is None or not self._server.sockets:
|
|
return self._port
|
|
return int(self._server.sockets[0].getsockname()[1])
|
|
|
|
def clients(self) -> list[WsClient]:
|
|
return list(self._clients)
|
|
|
|
async def start(self) -> None:
|
|
handler = self._handle_websocket if _NEW_WEBSOCKETS_API else self._handle_legacy_websocket
|
|
self._server = await ws_serve(handler, self._host, self._port)
|
|
|
|
async def stop(self) -> None:
|
|
if self._server is not None:
|
|
self._server.close()
|
|
await self._server.wait_closed()
|
|
self._server = None
|
|
clients = list(self._clients)
|
|
self._clients.clear()
|
|
for client in clients:
|
|
await client.close()
|
|
|
|
async def broadcast(self, message: Message) -> None:
|
|
for client in self.clients():
|
|
await client.communicator.send(message)
|
|
|
|
async def _handle_websocket(self, ws: Any) -> None:
|
|
if self._new_client is not None:
|
|
client = self._new_client(ws)
|
|
else:
|
|
if self._parser_map is None:
|
|
raise ValueError("parser_map is required when new_client is not provided")
|
|
client = WsClient(ws, self._interval_sec, self._wait_sec, self._parser_map)
|
|
|
|
self._clients.append(client)
|
|
client.add_disconnect_listener(lambda c: self._remove_client(c))
|
|
result = self.on_client_connected(client)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
await client.wait_closed()
|
|
|
|
def _remove_client(self, client: WsClient) -> None:
|
|
if client in self._clients:
|
|
self._clients.remove(client)
|
|
|
|
async def _handle_legacy_websocket(self, ws: Any, _path: str) -> None:
|
|
await self._handle_websocket(ws)
|