From 6b6f20a032eb31d747366c8bc270ca81731eaf6a Mon Sep 17 00:00:00 2001 From: Tom van der Lee Date: Fri, 4 Sep 2026 17:09:24 +0200 Subject: Fixed pre-commit issues --- proxy/websockets.py | 48 +++++++++++++++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 17 deletions(-) (limited to 'proxy/websockets.py') diff --git a/proxy/websockets.py b/proxy/websockets.py index ec65e6d..a09e261 100644 --- a/proxy/websockets.py +++ b/proxy/websockets.py @@ -9,17 +9,22 @@ from uuid import uuid4 from starlette.endpoints import WebSocketEndpoint from starlette.websockets import WebSocket -from proxy.utils import get_path_with_query_string from proxy.queue import ProxyQueue -from ttun_server.types import WebsocketMessage, WebsocketMessageType, WebsocketConnectData, WebsocketMessageData, \ - WebsocketDisconnectData +from proxy.utils import get_path_with_query_string +from ttun_server.types import ( + WebsocketConnectData, + WebsocketDisconnectData, + WebsocketMessage, + WebsocketMessageData, + WebsocketMessageType, +) logger = logging.getLogger(__name__) -logger.setLevel('DEBUG') +logger.setLevel("DEBUG") class WebsocketProxy(WebSocketEndpoint): - encoding = 'json' + encoding = "json" websocket_listen_task = None def __init__(self, *args, **kwargs): @@ -28,16 +33,20 @@ class WebsocketProxy(WebSocketEndpoint): @asynccontextmanager async def proxy(self, websocket: WebSocket, message: WebsocketMessage): - [subdomain, *_] = websocket.url.hostname.split('.') + [subdomain, *_] = websocket.url.hostname.split(".") - expect_ack = WebsocketMessageType(message['type']) == WebsocketMessageType.connect + expect_ack = ( + WebsocketMessageType(message["type"]) == WebsocketMessageType.connect + ) try: request_queue = await ProxyQueue.get_for_identifier(subdomain) await request_queue.enqueue(message) if expect_ack: - response_queue = await ProxyQueue.create_for_identifier(message["identifier"]) + response_queue = await ProxyQueue.create_for_identifier( + message["identifier"] + ) yield await response_queue.dequeue() await response_queue.delete() else: @@ -51,7 +60,9 @@ class WebsocketProxy(WebSocketEndpoint): while True: message: WebsocketMessage = await response_queue.dequeue() logger.debug(message) - await websocket.send_text(b64decode(message['payload']['body'].encode()).decode()) + await websocket.send_text( + b64decode(message["payload"]["body"].encode()).decode() + ) async def on_connect(self, websocket: WebSocket) -> None: message = WebsocketMessage( @@ -60,18 +71,21 @@ class WebsocketProxy(WebSocketEndpoint): payload=WebsocketConnectData( path=get_path_with_query_string(websocket), headers=[ - (k.decode(), v.decode()) - for k, v - in websocket.scope['headers'] + (k.decode(), v.decode()) for k, v in websocket.scope["headers"] ], - ) + ), ) async with self.proxy(websocket, message) as m: - if m is not None and WebsocketMessageType(m['type']) == WebsocketMessageType.ack: + if ( + m is not None + and WebsocketMessageType(m["type"]) == WebsocketMessageType.ack + ): await super().on_connect(websocket) - self.websocket_listen_task = asyncio.create_task(self.listen_for_messages(websocket)) + self.websocket_listen_task = asyncio.create_task( + self.listen_for_messages(websocket) + ) def callback(*args, **kwargs): self.websocket_listen_task = None @@ -92,7 +106,7 @@ class WebsocketProxy(WebSocketEndpoint): identifier=self.id, payload=WebsocketMessageData( body=b64encode(data_bytes).decode(), - ) + ), ) async with self.proxy(websocket, message): @@ -104,7 +118,7 @@ class WebsocketProxy(WebSocketEndpoint): identifier=self.id, payload=WebsocketDisconnectData( close_code=close_code, - ) + ), ) async with self.proxy(websocket, message): -- cgit v1.3