From 8733fb1faee89b97c78a7e9df4454f3523cea0d0 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 2 Nov 2022 07:23:14 +0100 Subject: [PATCH] Add type annotation to `wsproto_impl.py` --- setup.cfg | 1 + uvicorn/protocols/websockets/wsproto_impl.py | 81 +++++++++++++------- 2 files changed, 54 insertions(+), 28 deletions(-) diff --git a/setup.cfg b/setup.cfg index ac52be52e2..3a7c64ff53 100644 --- a/setup.cfg +++ b/setup.cfg @@ -37,6 +37,7 @@ files = uvicorn/protocols/http/__init__.py, uvicorn/protocols/websockets/__init__.py, uvicorn/protocols/websockets/websockets_impl.py, + uvicorn/protocols/websockets/wsproto_impl.py, uvicorn/protocols/http/h11_impl.py, uvicorn/protocols/http/httptools_impl.py, tests/middleware/test_wsgi.py, diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index a97766ff56..4f48714471 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -1,5 +1,6 @@ import asyncio import logging +import typing from urllib.parse import unquote import h11 @@ -9,6 +10,7 @@ from wsproto.extensions import PerMessageDeflate from wsproto.utilities import RemoteProtocolError +from uvicorn.config import Config from uvicorn.logging import TRACE_LOG_LEVEL from uvicorn.protocols.utils import ( get_local_addr, @@ -16,10 +18,28 @@ get_remote_addr, is_ssl, ) +from uvicorn.server import ServerState + +if typing.TYPE_CHECKING: + from asgiref.typing import ( + ASGISendEvent, + WebSocketAcceptEvent, + WebSocketCloseEvent, + WebSocketConnectEvent, + WebSocketDisconnectEvent, + WebSocketReceiveEvent, + WebSocketScope, + WebSocketSendEvent, + ) class WSProtocol(asyncio.Protocol): - def __init__(self, config, server_state, _loop=None): + def __init__( + self, + config: Config, + server_state: ServerState, + _loop: typing.Optional[asyncio.AbstractEventLoop] = None, + ) -> None: if not config.loaded: config.load() @@ -35,14 +55,14 @@ def __init__(self, config, server_state, _loop=None): self.default_headers = server_state.default_headers # Connection state - self.transport = None - self.server = None - self.client = None - self.scheme = None + self.transport: asyncio.Transport = None # type: ignore[assignment] + self.server: typing.Optional[typing.Tuple[str, int]] = None + self.client: typing.Optional[typing.Tuple[str, int]] = None + self.scheme: typing.Literal["wss", "ws"] = None # type: ignore[assignment] # WebSocket state self.connect_event = None - self.queue = asyncio.Queue() + self.queue: asyncio.Queue[WebSocketReceiveEvent] = asyncio.Queue() self.handshake_complete = False self.close_sent = False @@ -58,7 +78,9 @@ def __init__(self, config, server_state, _loop=None): # Protocol interface - def connection_made(self, transport): + def connection_made( # type: ignore[override] + self, transport: asyncio.Transport + ) -> None: self.connections.add(self) self.transport = transport self.server = get_local_addr(transport) @@ -66,33 +88,34 @@ def connection_made(self, transport): self.scheme = "wss" if is_ssl(transport) else "ws" if self.logger.level <= TRACE_LOG_LEVEL: - prefix = "%s:%d - " % tuple(self.client) if self.client else "" + prefix = "%s:%d - " % self.client if self.client else "" self.logger.log(TRACE_LOG_LEVEL, "%sWebSocket connection made", prefix) - def connection_lost(self, exc): + def connection_lost(self, exc: typing.Optional[Exception]) -> None: self.queue.put_nowait({"type": "websocket.disconnect"}) self.connections.remove(self) if self.logger.level <= TRACE_LOG_LEVEL: - prefix = "%s:%d - " % tuple(self.client) if self.client else "" + prefix = "%s:%d - " % self.client if self.client else "" self.logger.log(TRACE_LOG_LEVEL, "%sWebSocket connection lost", prefix) if exc is None: self.transport.close() - def eof_received(self): + def eof_received(self) -> None: pass - def data_received(self, data): + def data_received(self, data: bytes) -> None: try: self.conn.receive_data(data) except RemoteProtocolError as err: - self.transport.write(self.conn.send(err.event_hint)) + # TODO: Remove `type: ignore` when wsproto fixes the type annotation. + self.transport.write(self.conn.send(err.event_hint)) # type: ignore[arg-type] self.transport.close() else: self.handle_events() - def handle_events(self): + def handle_events(self) -> None: for event in self.conn.events(): if isinstance(event, events.Request): self.handle_connect(event) @@ -109,19 +132,19 @@ def handle_events(self): elif isinstance(event, events.Ping): self.handle_ping(event) - def pause_writing(self): + def pause_writing(self) -> None: """ Called by the transport when the write buffer exceeds the high water mark. """ self.writable.clear() - def resume_writing(self): + def resume_writing(self) -> None: """ Called by the transport when the write buffer drops below the low water mark. """ self.writable.set() - def shutdown(self): + def shutdown(self) -> None: if self.handshake_complete: self.queue.put_nowait({"type": "websocket.disconnect", "code": 1012}) output = self.conn.send(wsproto.events.CloseConnection(code=1012)) @@ -130,12 +153,12 @@ def shutdown(self): self.send_500_response() self.transport.close() - def on_task_complete(self, task): + def on_task_complete(self, task: asyncio.Task) -> None: self.tasks.discard(task) # Event handlers - def handle_connect(self, event): + def handle_connect(self, event: events.Request) -> None: self.connect_event = event headers = [(b"host", event.host.encode())] headers += [(key.lower(), value) for key, value in event.extra_headers] @@ -159,7 +182,9 @@ def handle_connect(self, event): task.add_done_callback(self.on_task_complete) self.tasks.add(task) - def handle_no_connect(self, event): + def handle_no_connect( + self, event: typing.Union[events.RejectData, events.RejectConnection] + ) -> None: headers = [ (b"content-type", b"text/plain; charset=utf-8"), (b"connection", b"close"), @@ -173,7 +198,7 @@ def handle_no_connect(self, event): self.transport.write(output) self.transport.close() - def handle_text(self, event): + def handle_text(self, event: events.TextMessage) -> None: self.text += event.data if event.message_finished: self.queue.put_nowait({"type": "websocket.receive", "text": self.text}) @@ -182,7 +207,7 @@ def handle_text(self, event): self.read_paused = True self.transport.pause_reading() - def handle_bytes(self, event): + def handle_bytes(self, event: events.BytesMessage) -> None: self.bytes += event.data # todo: we may want to guard the size of self.bytes and self.text if event.message_finished: @@ -192,16 +217,16 @@ def handle_bytes(self, event): self.read_paused = True self.transport.pause_reading() - def handle_close(self, event): + def handle_close(self, event: events.CloseConnection) -> None: if self.conn.state == ConnectionState.REMOTE_CLOSING: self.transport.write(self.conn.send(event.response())) self.queue.put_nowait({"type": "websocket.disconnect", "code": event.code}) self.transport.close() - def handle_ping(self, event): + def handle_ping(self, event: events.Ping) -> None: self.transport.write(self.conn.send(event.response())) - def send_500_response(self): + def send_500_response(self) -> None: headers = [ (b"content-type", b"text/plain; charset=utf-8"), (b"connection", b"close"), @@ -219,7 +244,7 @@ def send_500_response(self): output += self.conn.send(msg) self.transport.write(output) - async def run_asgi(self): + async def run_asgi(self) -> None: try: result = await self.app(self.scope, self.receive, self.send) except BaseException: @@ -238,7 +263,7 @@ async def run_asgi(self): self.logger.error(msg, result) self.transport.close() - async def send(self, message): + async def send(self, message: WebSocketSendEvent) -> None: await self.writable.wait() message_type = message["type"] @@ -319,7 +344,7 @@ async def send(self, message): msg = "Unexpected ASGI message '%s', after sending 'websocket.close'." raise RuntimeError(msg % message_type) - async def receive(self): + async def receive(self) -> WebSocketReceiveEvent: message = await self.queue.get() if self.read_paused and self.queue.empty(): self.read_paused = False