Add utility to wait for several events.

Indeed, trio has no equivalent of asyncio.wait.

Author: aaugustin

src/websockets/trio/utils.py

@@ -0,0 +1,28 @@
+import trio
+
+
+__all__ = ["wait_for_any_event"]
+
+
+# Based on https://trio.readthedocs.io/en/stable/reference-core.html#custom-supervisors
+
+
+async def jockey(event: trio.Event, cancel_scope: trio.CancelScope) -> None:
+    await event.wait()
+    cancel_scope.cancel()
+
+
+async def wait_for_any_event(*events: trio.Event) -> None:
+    """
+    Wait for any of the given events to be set.
+
+    Args:
+        *events: The events to wait for.
+
+    """
+    if not events:
+        raise ValueError("no events provided")
+
+    async with trio.open_nursery() as nursery:
+        for event in events:
+            nursery.start_soon(jockey, event, nursery.cancel_scope)

tests/trio/test_utils.py

@@ -0,0 +1,28 @@
+import trio.testing
+
+from websockets.trio.utils import *
+
+from .utils import IsolatedTrioTestCase
+
+
+class UtilsTests(IsolatedTrioTestCase):
+    async def test_wait_for_any_event(self):
+        event1 = trio.Event()
+        event2 = trio.Event()
+        done = trio.Event()
+
+        async def waiter():
+            await wait_for_any_event(event1, event2)
+            done.set()
+
+        self.nursery.start_soon(waiter)
+        await trio.testing.wait_all_tasks_blocked()
+        self.assertFalse(done.is_set())
+
+        event1.set()
+        await trio.testing.wait_all_tasks_blocked()
+        self.assertTrue(done.is_set())
+
+    async def test_wait_for_any_event_no_events(self):
+        with self.assertRaises(ValueError):
+            await wait_for_any_event()