WebSockets¶

defurl_for(self, name: str, /, **path_params: Any) -> URL:
 url_path_provider: Router | Starlette | None = self.scope.get("router") or self.scope.get("app")
 if url_path_provider is None:
 raise RuntimeError("The `url_for` method can only be used inside a Starlette application or with a router.")
 url_path = url_path_provider.url_path_for(name, **path_params)
 return url_path.make_absolute_url(base_url=self.base_url)

async defreceive(self) -> Message:
"""
 Receive ASGI websocket messages, ensuring valid state transitions.
 """
 if self.client_state == WebSocketState.CONNECTING:
 message = await self._receive()
 message_type = message["type"]
 if message_type != "websocket.connect":
 raise RuntimeError(f'Expected ASGI message "websocket.connect", but got {message_type!r}')
 self.client_state = WebSocketState.CONNECTED
 return message
 elif self.client_state == WebSocketState.CONNECTED:
 message = await self._receive()
 message_type = message["type"]
 if message_type not in {"websocket.receive", "websocket.disconnect"}:
 raise RuntimeError(
 f'Expected ASGI message "websocket.receive" or "websocket.disconnect", but got {message_type!r}'
 )
 if message_type == "websocket.disconnect":
 self.client_state = WebSocketState.DISCONNECTED
 return message
 else:
 raise RuntimeError('Cannot call "receive" once a disconnect message has been received.')

async defsend(self, message: Message) -> None:
"""
 Send ASGI websocket messages, ensuring valid state transitions.
 """
 if self.application_state == WebSocketState.CONNECTING:
 message_type = message["type"]
 if message_type not in {"websocket.accept", "websocket.close", "websocket.http.response.start"}:
 raise RuntimeError(
 'Expected ASGI message "websocket.accept", "websocket.close" or "websocket.http.response.start", '
 f"but got {message_type!r}"
 )
 if message_type == "websocket.close":
 self.application_state = WebSocketState.DISCONNECTED
 elif message_type == "websocket.http.response.start":
 self.application_state = WebSocketState.RESPONSE
 else:
 self.application_state = WebSocketState.CONNECTED
 await self._send(message)
 elif self.application_state == WebSocketState.CONNECTED:
 message_type = message["type"]
 if message_type not in {"websocket.send", "websocket.close"}:
 raise RuntimeError(
 f'Expected ASGI message "websocket.send" or "websocket.close", but got {message_type!r}'
 )
 if message_type == "websocket.close":
 self.application_state = WebSocketState.DISCONNECTED
 try:
 await self._send(message)
 except OSError:
 self.application_state = WebSocketState.DISCONNECTED
 raise WebSocketDisconnect(code=1006)
 elif self.application_state == WebSocketState.RESPONSE:
 message_type = message["type"]
 if message_type != "websocket.http.response.body":
 raise RuntimeError(f'Expected ASGI message "websocket.http.response.body", but got {message_type!r}')
 if not message.get("more_body", False):
 self.application_state = WebSocketState.DISCONNECTED
 await self._send(message)
 else:
 raise RuntimeError('Cannot call "send" once a close message has been sent.')

async defaccept(
 self,
 subprotocol: str | None = None,
 headers: Iterable[tuple[bytes, bytes]] | None = None,
) -> None:
 headers = headers or []

 if self.client_state == WebSocketState.CONNECTING: # pragma: no branch
 # If we haven't yet seen the 'connect' message, then wait for it first.
 await self.receive()
 await self.send({"type": "websocket.accept", "subprotocol": subprotocol, "headers": headers})

async defreceive_text(self) -> str:
 if self.application_state != WebSocketState.CONNECTED:
 raise RuntimeError('WebSocket is not connected. Need to call "accept" first.')
 message = await self.receive()
 self._raise_on_disconnect(message)
 return cast(str, message["text"])

async defreceive_bytes(self) -> bytes:
 if self.application_state != WebSocketState.CONNECTED:
 raise RuntimeError('WebSocket is not connected. Need to call "accept" first.')
 message = await self.receive()
 self._raise_on_disconnect(message)
 return cast(bytes, message["bytes"])

async defreceive_json(self, mode: str = "text") -> Any:
 if mode not in {"text", "binary"}:
 raise RuntimeError('The "mode" argument should be "text" or "binary".')
 if self.application_state != WebSocketState.CONNECTED:
 raise RuntimeError('WebSocket is not connected. Need to call "accept" first.')
 message = await self.receive()
 self._raise_on_disconnect(message)

 if mode == "text":
 text = message["text"]
 else:
 text = message["bytes"].decode("utf-8")
 return json.loads(text)

async defiter_text(self) -> AsyncIterator[str]:
 try:
 while True:
 yield await self.receive_text()
 except WebSocketDisconnect:
 pass

async defiter_bytes(self) -> AsyncIterator[bytes]:
 try:
 while True:
 yield await self.receive_bytes()
 except WebSocketDisconnect:
 pass

async defiter_json(self) -> AsyncIterator[Any]:
 try:
 while True:
 yield await self.receive_json()
 except WebSocketDisconnect:
 pass

async defsend_text(self, data: str) -> None:
 await self.send({"type": "websocket.send", "text": data})

async defsend_bytes(self, data: bytes) -> None:
 await self.send({"type": "websocket.send", "bytes": data})

async defsend_json(self, data: Any, mode: str = "text") -> None:
 if mode not in {"text", "binary"}:
 raise RuntimeError('The "mode" argument should be "text" or "binary".')
 text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
 if mode == "text":
 await self.send({"type": "websocket.send", "text": text})
 else:
 await self.send({"type": "websocket.send", "bytes": text.encode("utf-8")})

async defclose(self, code: int = 1000, reason: str | None = None) -> None:
 await self.send({"type": "websocket.close", "code": code, "reason": reason or ""})