add send and close event

encode · Nov 20, 2022 · 94728bd · 94728bd
1 parent 9cee0b4
commit 94728bd
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py
@@ -24,10 +24,12 @@
     from asgiref.typing import (
         ASGISendEvent,
         WebSocketAcceptEvent,
+        WebSocketCloseEvent,
         WebSocketConnectEvent,
         WebSocketDisconnectEvent,
         WebSocketReceiveEvent,
         WebSocketScope,
+        WebSocketSendEvent,
     )
 
     WebSocketEvent = typing.Union[
@@ -308,14 +310,18 @@ async def send(self, message: "ASGISendEvent") -> None:
 
         elif not self.close_sent:
             if message_type == "websocket.send":
+                message = typing.cast("WebSocketSendEvent", message)
                 bytes_data = message.get("bytes")
                 text_data = message.get("text")
                 data = text_data if bytes_data is None else bytes_data
-                output = self.conn.send(wsproto.events.Message(data=data))
+                output = self.conn.send(
+                    wsproto.events.Message(data=data)  # type: ignore[type-var]
+                )
                 if not self.transport.is_closing():
                     self.transport.write(output)
 
             elif message_type == "websocket.close":
+                message = typing.cast("WebSocketCloseEvent", message)
                 self.close_sent = True
                 code = message.get("code", 1000)
                 reason = message.get("reason", "") or ""