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
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from google.protobuf.message import DecodeError
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
try:
|
|
from websockets.asyncio.client import connect as ws_connect
|
|
except ImportError: # websockets 12.x legacy API
|
|
from websockets import connect as ws_connect
|
|
|
|
from toki_socket.base_client import BaseClient
|
|
from toki_socket.communicator import ParserMap, type_name_of
|
|
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
|
|
|
|
|
class WsClient(BaseClient):
|
|
def __init__(
|
|
self,
|
|
ws: Any,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> None:
|
|
super().__init__(interval_sec, wait_sec, self._do_close_impl)
|
|
self._ws = ws
|
|
self._init_base()
|
|
self.communicator.initialize(self, parser_map)
|
|
self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
|
|
self.communicator.add_listener(
|
|
type_name_of(HeartBeat),
|
|
lambda _: asyncio.create_task(self._on_heartbeat()),
|
|
)
|
|
self._read_task = asyncio.create_task(self._read_loop())
|
|
asyncio.create_task(self._send_heartbeat())
|
|
|
|
async def write_packet(self, base: PacketBase) -> None:
|
|
await self._ws.send(base.SerializeToString())
|
|
|
|
async def wait_closed(self) -> None:
|
|
await self.communicator.wait_closed()
|
|
|
|
async def _do_close_impl(self) -> None:
|
|
with suppress(Exception):
|
|
await self._ws.close()
|
|
|
|
async def _read_loop(self) -> None:
|
|
try:
|
|
async for message in self._ws:
|
|
if not isinstance(message, bytes):
|
|
continue
|
|
base = PacketBase()
|
|
base.ParseFromString(message)
|
|
self.communicator.on_received_data(
|
|
base.typeName,
|
|
base.data,
|
|
base.nonce,
|
|
base.responseNonce,
|
|
)
|
|
asyncio.create_task(self._send_heartbeat())
|
|
except (ConnectionClosed, ConnectionError, OSError, DecodeError):
|
|
pass
|
|
finally:
|
|
await self._on_disconnected()
|
|
|
|
|
|
async def connect_ws(
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> WsClient:
|
|
uri = f"ws://{host}:{port}{path}"
|
|
ws = await ws_connect(uri)
|
|
return WsClient(ws, interval_sec, wait_sec, parser_map)
|