- add subscription via MessageDispatcher

ws/client/ws_client.py

@@ -8,9 +8,9 @@ from ws.connection import IConnection
 
 
 class WsClient:
-	def __init__(self, listener: IConnection):
+	def __init__(self, listener: IConnection, loop=None):
 		self.listener = listener
-		self.loop = asyncio.new_event_loop()
+		self.loop = loop
 		self.stop = None
 		self.bg_thread = None
 
@@ -44,7 +44,7 @@ class WsClient:
 			return
 
 		print("handler: got connection to {}".format(websocket.remote_address))
-		self.listener.on_connect(self.loop)
+		await self.listener.on_connect()
 		try:
 			path = "/"
 			consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path), loop=self.loop)
@@ -55,7 +55,7 @@ class WsClient:
 		finally:
 			print("handler: lost connection from {}".format(websocket.remote_address))
 			await websocket.close()
-			self.listener.on_disconnect()
+			await self.listener.on_disconnect()
 
 	async def handler_recv(self, websocket, path):
 		while True:

ws/message.py

@@ -45,9 +45,11 @@ class Value:
 
 
 class MessageDispatcher(IConnection):
-	def __init__(self):
+	def __init__(self, auto_subscribe=False, loop=None):
 		self.msg_handlers: MsgIo = []
 		self.state = None
+		self.state = asyncio.Queue(loop=loop)
+		self.auto_subscribe = auto_subscribe
 
 	def msgio_get(self, key):
 		obj = MsgIo(key, self.send)
@@ -56,10 +58,14 @@ class MessageDispatcher(IConnection):
 			return obj
 		return None
 
-	def on_connect(self, loop):
-		self.state = asyncio.Queue(loop=loop)
+	async def on_connect(self):
+		if self.auto_subscribe:
+			for handler in self.msg_handlers:
+				key = handler.get_key()
+				print("Would subscribe {}".format(key))
+				await self.send({'+': key})
 
-	def on_disconnect(self):
+	async def on_disconnect(self):
 		pass
 
 	async def on_recv(self, data):

ws/server/ws_server.py

@@ -5,10 +5,13 @@ import abc
 
 class WsServer:
 	def __init__(self, loop=None):
-		if loop is None:
-			self.loop = asyncio.get_event_loop()
+		self.loop = loop
+
+	def run_forever(self):
+		if self.loop is None:
+			asyncio.get_event_loop().run_forever()
 		else:
-			self.loop = loop
+			self.loop.run_forever()
 
 	async def listen(self, host, port):
 		return await websockets.serve(self.handler, host, port, loop=self.loop)

ws/server/ws_server_multi_user.py

@@ -21,7 +21,6 @@ class WsServerMultiUser(WsServer):
 		self.listener = listener
 		self.USERS = UserSet()
 		self.global_state = {}
-		self.listener.on_connect(self.loop)
 
 	async def notify_state(self, data):
 		if self.USERS:  # asyncio.wait doesn't accept an empty list