Skip to content

Commit f2036db

Browse files
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent 8c816f0 commit f2036db

File tree

4 files changed

+73
-64
lines changed

4 files changed

+73
-64
lines changed

integration_tests/base_routes.py

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from typing import Optional
88

99
from integration_tests.subroutes import di_subrouter, sub_router
10-
from robyn import Headers, Request, Response, Robyn, SSEMessage, SSEResponse, WebSocketConnector, WebSocketDisconnect, jsonify, serve_file, serve_html
10+
from robyn import Headers, Request, Response, Robyn, SSEMessage, SSEResponse, WebSocketDisconnect, jsonify, serve_file, serve_html
1111
from robyn.authentication import AuthenticationHandler, BearerGetter, Identity
1212
from robyn.robyn import QueryParams, Url
1313
from robyn.templating import JinjaTemplate
@@ -28,14 +28,14 @@
2828
@app.websocket("/web_socket")
2929
async def websocket_endpoint(websocket):
3030
await websocket.accept()
31-
31+
3232
try:
3333
while True:
3434
msg = await websocket.receive_text()
3535
websocket_id = websocket._connector.id
3636
global websocket_state
3737
state = websocket_state[websocket_id]
38-
38+
3939
if state == 0:
4040
await websocket._connector.async_broadcast("This is a broadcast message")
4141
await websocket.send_text("This is a message to self")
@@ -50,9 +50,9 @@ async def websocket_endpoint(websocket):
5050
await websocket.send_text("Connection closed")
5151
await websocket.close()
5252
break
53-
53+
5454
websocket_state[websocket_id] = (state + 1) % 4
55-
55+
5656
except WebSocketDisconnect:
5757
pass
5858

@@ -62,7 +62,7 @@ def websocket_on_connect(websocket):
6262
return "Hello world, from ws"
6363

6464

65-
@websocket_endpoint.on_close
65+
@websocket_endpoint.on_close
6666
def websocket_on_close(websocket):
6767
return "GoodBye world, from ws"
6868

@@ -71,25 +71,25 @@ def websocket_on_close(websocket):
7171
@app.websocket("/web_socket_json")
7272
async def json_websocket_endpoint(websocket):
7373
await websocket.accept()
74-
74+
7575
try:
7676
while True:
7777
msg = await websocket.receive_text()
7878
websocket_id = websocket._connector.id
7979
response = {"ws_id": websocket_id, "resp": "", "msg": msg}
8080
global websocket_state
8181
state = websocket_state[websocket_id]
82-
82+
8383
if state == 0:
8484
response["resp"] = "Whaaat??"
8585
elif state == 1:
8686
response["resp"] = "Whooo??"
8787
elif state == 2:
8888
response["resp"] = "*chika* *chika* Slim Shady."
89-
89+
9090
websocket_state[websocket_id] = (state + 1) % 3
9191
await websocket.send_json(response)
92-
92+
9393
except WebSocketDisconnect:
9494
pass
9595

@@ -108,13 +108,13 @@ def json_websocket_on_close(websocket):
108108
@app.websocket("/web_socket_di")
109109
async def di_websocket_endpoint(websocket):
110110
await websocket.accept()
111-
111+
112112
try:
113113
while True:
114114
await websocket.receive_text()
115115
# Just echo back an empty response for DI test
116116
await websocket.send_text("")
117-
117+
118118
except WebSocketDisconnect:
119119
pass
120120

robyn/__init__.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,6 @@
3030
__version__ = get_version()
3131

3232

33-
34-
3533
def _normalize_endpoint(endpoint: str) -> str:
3634
"""
3735
Normalize an endpoint to ensure consistent routing.
@@ -286,21 +284,21 @@ def websocket(self, endpoint: str):
286284
"""
287285
Modern WebSocket decorator that accepts a single handler function.
288286
The handler function receives a WebSocket object and can optionally have on_connect and on_close callbacks.
289-
287+
290288
Usage:
291289
@app.websocket("/ws")
292290
async def websocket_endpoint(websocket):
293291
await websocket.accept()
294292
while True:
295293
data = await websocket.receive_text()
296294
await websocket.send_text(f"Echo: {data}")
297-
295+
298296
# With optional callbacks:
299297
@websocket_endpoint.on_connect
300298
async def on_connect(websocket):
301299
await websocket.send_text("Connected!")
302-
303-
@websocket_endpoint.on_close
300+
301+
@websocket_endpoint.on_close
304302
async def on_close(websocket):
305303
print("Disconnected")
306304
"""
@@ -711,15 +709,15 @@ def websocket(self, endpoint: str):
711709
"""
712710
Modern WebSocket decorator for SubRouter that accepts a single handler function.
713711
Works the same as the main Robyn websocket decorator but with prefix support.
714-
712+
715713
Usage:
716714
@subrouter.websocket("/ws")
717715
async def websocket_endpoint(websocket):
718716
await websocket.accept()
719717
while True:
720718
data = await websocket.receive_text()
721719
await websocket.send_text(f"Echo: {data}")
722-
720+
723721
# With optional callbacks:
724722
@websocket_endpoint.on_connect
725723
async def on_connect(websocket):

robyn/ws.py

Lines changed: 45 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
class 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

Comments
 (0)