From a79c408fcdbbb76ad2fffe9bcfce334f0b8889f7 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Mon, 19 Jun 2023 10:52:38 +0200 Subject: [PATCH] Type get_room as possible coroutine --- ypy_websocket/websocket_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ypy_websocket/websocket_server.py b/ypy_websocket/websocket_server.py index c9717fd..a586e88 100644 --- a/ypy_websocket/websocket_server.py +++ b/ypy_websocket/websocket_server.py @@ -3,6 +3,7 @@ from contextlib import AsyncExitStack from inspect import isawaitable from logging import Logger, getLogger +from typing import Awaitable, cast from anyio import TASK_STATUS_IGNORED, Event, create_task_group from anyio.abc import TaskGroup, TaskStatus @@ -56,7 +57,7 @@ def started(self) -> Event: self._started = Event() return self._started - def get_room(self, name: str) -> YRoom: + def get_room(self, name: str) -> YRoom | Awaitable[YRoom]: """Get or create a room with the given name. Arguments: @@ -130,6 +131,7 @@ async def _serve(self, websocket, *, task_status: TaskStatus[None] = TASK_STATUS room = self.get_room(websocket.path) if isawaitable(room): room = await room + room = cast(YRoom, room) if not room.started.is_set(): tg.start_soon(room.start) await room.started.wait()