Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tqxrk8uj4M3w3d3eXm3xK8
41 lines
956 B
Python
41 lines
956 B
Python
import asyncio
|
|
import websockets
|
|
import abc
|
|
|
|
|
|
class WsServer:
|
|
def __init__(self, loop=None):
|
|
self.loop = loop
|
|
|
|
def run_forever(self):
|
|
if self.loop is None:
|
|
asyncio.get_event_loop().run_forever()
|
|
else:
|
|
self.loop.run_forever()
|
|
|
|
async def listen(self, host, port):
|
|
return await websockets.serve(self.handler, host, port)
|
|
|
|
def disconnect(self):
|
|
pass
|
|
|
|
async def handler(self, websocket, path):
|
|
await self.register(websocket)
|
|
try:
|
|
consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path))
|
|
producer_task = asyncio.ensure_future(self.handler_send(websocket, path))
|
|
done, pending = await asyncio.wait([consumer_task, producer_task], return_when=asyncio.FIRST_COMPLETED,)
|
|
for task in pending:
|
|
task.cancel()
|
|
finally:
|
|
await self.unregister(websocket)
|
|
|
|
@abc.abstractmethod
|
|
async def handler_recv(self, websocket, path):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def handler_send(self, websocket, path):
|
|
pass
|
|
|