11import asyncio
22import json
33import logging
4- from typing import List , Set
4+ from typing import List , TYPE_CHECKING , Set
55import websockets
66
7+ if TYPE_CHECKING :
8+ from event_handlers .base_event_handler import EventHandler
9+
710
811class BrowserWebsocketServer :
912 """
@@ -31,15 +34,15 @@ def __init__(self):
3134 so we can use them to send outbound messages from this plugin to the
3235 extension.
3336 """
34- self ._ws_clients : Set [websockets .WebSocketServerProtocol ] = set ()
37+ self ._ws_clients : Set [websockets .ServerConnection ] = set ()
3538
3639 """
3740 Any EventHandlers registered to receive inbound events from the browser extension.
3841 """
3942 self ._handlers : List ["EventHandler" ] = []
4043
41- def start (self , hostname : str , port : int ) -> None :
42- return websockets .serve (self ._message_receive_loop , hostname , port )
44+ async def start (self , hostname : str , port : int ) -> websockets . Server :
45+ return await websockets .serve (self ._message_receive_loop , hostname , port )
4346
4447 async def send_to_clients (self , message : str ) -> None :
4548 """
@@ -66,16 +69,16 @@ def register_event_handler(self, handler: "EventHandler") -> None:
6669 def num_connected_clients (self ) -> int :
6770 return len (self ._ws_clients )
6871
69- def _register_client (self , ws : websockets .WebSocketServerProtocol ) -> None :
72+ def _register_client (self , ws : websockets .ServerConnection ) -> None :
7073 self ._ws_clients .add (ws )
7174 self ._logger .info (
7275 (f"{ ws .remote_address } has connected to our browser websocket."
7376 f" We now have { len (self ._ws_clients )} active connection(s)." ))
7477
75- async def _unregister_client (self , ws : websockets .WebSocketServerProtocol ) -> None :
78+ async def _unregister_client (self , ws : websockets .ServerConnection ) -> None :
7679 try :
7780 await ws .close ()
78- except :
81+ except Exception :
7982 self ._logger .exception (
8083 "Exception while closing browser webocket connection." )
8184 if ws in self ._ws_clients :
@@ -84,7 +87,7 @@ async def _unregister_client(self, ws: websockets.WebSocketServerProtocol) -> No
8487 (f"{ ws .remote_address } has disconnected from our browser websocket."
8588 f" We now have { len (self ._ws_clients )} active connection(s) remaining." ))
8689
87- async def _message_receive_loop (self , ws : websockets .WebSocketServerProtocol , uri : str ) -> None :
90+ async def _message_receive_loop (self , ws : websockets .ServerConnection ) -> None :
8891 """
8992 Loop of waiting for and processing inbound websocket messages, until the
9093 connection dies. Each connection will create one of these coroutines.
@@ -93,9 +96,9 @@ async def _message_receive_loop(self, ws: websockets.WebSocketServerProtocol, ur
9396 try :
9497 async for message in ws :
9598 self ._logger .info (
96- f"Received inbound message from browser extension. Message: { message } " )
99+ f"Received inbound message from browser extension. Message: { str ( message ) } " )
97100 await self ._process_inbound_message (message )
98- except :
101+ except Exception :
99102 self ._logger .exception (
100103 "BrowserWebsocketServer encountered an exception while waiting for inbound messages." )
101104 finally :
@@ -105,24 +108,24 @@ async def _message_receive_loop(self, ws: websockets.WebSocketServerProtocol, ur
105108 for handler in self ._handlers :
106109 try :
107110 await handler .on_all_browsers_disconnected ()
108- except :
111+ except Exception :
109112 self ._logger .exception (
110113 "Connection mananger received an exception from EventHandler!" )
111114
112- async def _process_inbound_message (self , message : str ) -> None :
115+ async def _process_inbound_message (self , message : str | bytes ) -> None :
113116 """
114117 Process one individual inbound websocket message.
115118 """
116119 try :
117120 parsed_event = json .loads (message )
118- except :
121+ except Exception :
119122 self ._logger .exception (
120- f"Failed to parse browser websocket message as JSON. Message: { message } " )
123+ f"Failed to parse browser websocket message as JSON. Message: { str ( message ) } " )
121124 return
122125
123126 for handler in self ._handlers :
124127 try :
125128 await handler .on_browser_event (parsed_event )
126- except :
129+ except Exception :
127130 self ._logger .exception (
128131 "Connection mananger received an exception from EventHandler!" )
0 commit comments