43 lines
1.1 KiB
Python
43 lines
1.1 KiB
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, loop=self.loop)
|
|
|
|
def disconnect(self):
|
|
pass
|
|
|
|
async def handler(self, websocket, path):
|
|
print("handler: got connection from {}".format(websocket.remote_address))
|
|
await self.register(websocket)
|
|
try:
|
|
consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path), loop=self.loop)
|
|
producer_task = asyncio.ensure_future(self.handler_send(websocket, path), loop=self.loop)
|
|
done, pending = await asyncio.wait([consumer_task, producer_task], return_when=asyncio.FIRST_COMPLETED,)
|
|
for task in pending:
|
|
task.cancel()
|
|
finally:
|
|
print("handler: lost connection from {}".format(websocket.remote_address))
|
|
await self.unregister(websocket)
|
|
|
|
@abc.abstractmethod
|
|
async def handler_recv(self, websocket, path):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def handler_send(self, websocket, path):
|
|
pass
|
|
|