1515
1616class WebSocketDisconnect (Exception ):
1717 """Exception raised when a WebSocket connection is disconnected."""
18-
18+
1919 def __init__ (self , code : int = 1000 , reason : str = "" ):
2020 self .code = code
2121 self .reason = reason
@@ -27,28 +27,28 @@ class WebSocketAdapter:
2727 Adapter class that provides a modern WebSocket interface
2828 wrapping Robyn's WebSocketConnector for compatibility.
2929 """
30-
30+
3131 def __init__ (self , websocket_connector : WebSocketConnector , message : str = None ):
3232 self ._connector = websocket_connector
3333 self ._message = message
3434 self ._accepted = False
35-
35+
3636 async def accept (self ):
3737 """Accept the WebSocket connection (no-op in Robyn as it's auto-accepted)"""
3838 self ._accepted = True
39-
39+
4040 async def close (self , code : int = 1000 ):
4141 """Close the WebSocket connection"""
4242 self ._connector .close ()
43-
43+
4444 async def send_text (self , data : str ):
4545 """Send text data to the WebSocket"""
4646 await self ._connector .async_send_to (self ._connector .id , data )
47-
47+
4848 async def send_bytes (self , data : bytes ):
4949 """Send binary data to the WebSocket"""
50- await self ._connector .async_send_to (self ._connector .id , data .decode (' utf-8' ))
51-
50+ await self ._connector .async_send_to (self ._connector .id , data .decode (" utf-8" ))
51+
5252 async def receive_text (self ) -> str :
5353 """Receive text data from the WebSocket"""
5454 if self ._message is not None :
@@ -58,43 +58,45 @@ async def receive_text(self) -> str:
5858 # Note: In a real implementation, this would need to handle the message queue
5959 # For now, we return the current message if available
6060 return ""
61-
61+
6262 async def receive_bytes (self ) -> bytes :
6363 """Receive binary data from the WebSocket"""
6464 text = await self .receive_text ()
65- return text .encode (' utf-8' )
66-
65+ return text .encode (" utf-8" )
66+
6767 async def send_json (self , data ):
6868 """Send JSON data to the WebSocket"""
6969 import json
70+
7071 await self .send_text (json .dumps (data ))
71-
72+
7273 async def receive_json (self ):
7374 """Receive JSON data from the WebSocket"""
7475 import json
76+
7577 text = await self .receive_text ()
7678 return json .loads (text ) if text else None
77-
79+
7880 @property
7981 def query_params (self ):
8082 """Access query parameters"""
8183 return self ._connector .query_params
82-
84+
8385 @property
8486 def path_params (self ):
8587 """Access path parameters"""
86- return getattr (self ._connector , ' path_params' , {})
87-
88+ return getattr (self ._connector , " path_params" , {})
89+
8890 @property
8991 def headers (self ):
9092 """Access request headers"""
91- return getattr (self ._connector , ' headers' , {})
92-
93+ return getattr (self ._connector , " headers" , {})
94+
9395 @property
9496 def client (self ):
9597 """Client information"""
96- return getattr (self ._connector , ' client' , None )
97-
98+ return getattr (self ._connector , " client" , None )
99+
98100 @property
99101 def id (self ):
100102 """WebSocket connection ID"""
@@ -106,32 +108,34 @@ def create_websocket_decorator(app_instance):
106108 Factory function to create a websocket decorator for an app instance.
107109 This allows access to the app's dependencies and web_socket_router.
108110 """
111+
109112 def websocket (endpoint : str ):
110113 """
111114 Modern WebSocket decorator that accepts a single handler function.
112115 The handler function receives a WebSocket object and can optionally have on_connect and on_close callbacks.
113-
116+
114117 Usage:
115118 @app.websocket("/ws")
116119 async def websocket_endpoint(websocket):
117120 await websocket.accept()
118121 while True:
119122 data = await websocket.receive_text()
120123 await websocket.send_text(f"Echo: {data}")
121-
124+
122125 # With optional callbacks:
123126 @websocket_endpoint.on_connect
124127 async def on_connect(websocket):
125128 await websocket.send_text("Connected!")
126-
127- @websocket_endpoint.on_close
129+
130+ @websocket_endpoint.on_close
128131 async def on_close(websocket):
129132 print("Disconnected")
130133 """
134+
131135 def decorator (handler ):
132136 # Dictionary to store handlers for this WebSocket endpoint
133137 handlers = {}
134-
138+
135139 # Create the main message handler
136140 async def message_handler (websocket_connector , msg , * args , ** kwargs ):
137141 # Convert WebSocketConnector to modern WebSocket interface
@@ -148,78 +152,78 @@ async def message_handler(websocket_connector, msg, *args, **kwargs):
148152 if "connection closed" in str (e ).lower () or "websocket" in str (e ).lower ():
149153 return ""
150154 raise e
151-
155+
152156 # Create FunctionInfo for the message handler
153157 params = dict (inspect .signature (message_handler ).parameters )
154158 num_params = len (params )
155159 is_async = asyncio .iscoroutinefunction (message_handler )
156160 injected_dependencies = app_instance .dependencies .get_dependency_map (app_instance )
157-
161+
158162 # Filter dependencies to only include reserved parameters that exist in handler
159163 filtered_dependencies = {}
160164 if "global_dependencies" in params :
161165 filtered_dependencies ["global_dependencies" ] = injected_dependencies .get ("global_dependencies" , {})
162166 if "router_dependencies" in params :
163167 filtered_dependencies ["router_dependencies" ] = injected_dependencies .get ("router_dependencies" , {})
164-
168+
165169 handlers ["message" ] = FunctionInfo (message_handler , is_async , num_params , params , filtered_dependencies )
166-
170+
167171 # Add methods to the handler to allow attaching on_connect and on_close
168172 def add_on_connect (connect_handler ):
169173 def connect_wrapper (websocket_connector , * args , ** kwargs ):
170174 websocket_adapter = WebSocketAdapter (websocket_connector )
171175 if asyncio .iscoroutinefunction (connect_handler ):
172176 return asyncio .create_task (connect_handler (websocket_adapter ))
173177 return connect_handler (websocket_adapter )
174-
178+
175179 # Create FunctionInfo for connect handler
176180 connect_params = dict (inspect .signature (connect_handler ).parameters ) # Use original handler params, not wrapper
177181 connect_num_params = len (connect_params )
178182 connect_is_async = asyncio .iscoroutinefunction (connect_wrapper )
179-
183+
180184 # Filter dependencies for connect handler - only reserved parameters
181185 filtered_connect_deps = {}
182186 if "global_dependencies" in connect_params :
183187 filtered_connect_deps ["global_dependencies" ] = injected_dependencies .get ("global_dependencies" , {})
184188 if "router_dependencies" in connect_params :
185189 filtered_connect_deps ["router_dependencies" ] = injected_dependencies .get ("router_dependencies" , {})
186-
190+
187191 handlers ["connect" ] = FunctionInfo (connect_wrapper , connect_is_async , connect_num_params , connect_params , filtered_connect_deps )
188192 return connect_handler
189-
193+
190194 def add_on_close (close_handler ):
191195 def close_wrapper (websocket_connector , * args , ** kwargs ):
192196 websocket_adapter = WebSocketAdapter (websocket_connector )
193197 if asyncio .iscoroutinefunction (close_handler ):
194198 return asyncio .create_task (close_handler (websocket_adapter ))
195199 return close_handler (websocket_adapter )
196-
200+
197201 # Create FunctionInfo for close handler
198202 close_params = dict (inspect .signature (close_handler ).parameters ) # Use original handler params, not wrapper
199203 close_num_params = len (close_params )
200204 close_is_async = asyncio .iscoroutinefunction (close_wrapper )
201-
205+
202206 # Filter dependencies for close handler - only reserved parameters
203207 filtered_close_deps = {}
204208 if "global_dependencies" in close_params :
205209 filtered_close_deps ["global_dependencies" ] = injected_dependencies .get ("global_dependencies" , {})
206210 if "router_dependencies" in close_params :
207211 filtered_close_deps ["router_dependencies" ] = injected_dependencies .get ("router_dependencies" , {})
208-
212+
209213 handlers ["close" ] = FunctionInfo (close_wrapper , close_is_async , close_num_params , close_params , filtered_close_deps )
210214 return close_handler
211-
215+
212216 # Attach methods to the handler function
213217 handler .on_connect = add_on_connect
214218 handler .on_close = add_on_close
215219 handler ._ws_handlers = handlers # Store reference to handlers dict
216-
220+
217221 # Add the WebSocket to the router
218222 app_instance .add_web_socket (endpoint , handlers )
219223 return handler
220-
224+
221225 return decorator
222-
226+
223227 return websocket
224228
225229
@@ -255,4 +259,4 @@ def inner(handler):
255259
256260 return handler
257261
258- return inner
262+ return inner
0 commit comments