gen: accept any _Yieldable in WaitIterator constructor (issue #3001)

tornado/gen.py

@@ -350,23 +350,24 @@ class WaitIterator:
 
     _unfinished: dict[Future, int | str] = {}
 
-    def __init__(self, *args: Future, **kwargs: Future) -> None:
+    def __init__(self, *args: _Yieldable, **kwargs: _Yieldable) -> None:
         if args and kwargs:
             raise ValueError("You must provide args or kwargs, not both")
 
         if kwargs:
-            self._unfinished = {f: k for (k, f) in kwargs.items()}
-            futures: Sequence[Future] = list(kwargs.values())
+            self._unfinished = {convert_yielded(fut): k for (k, fut) in kwargs.items()}
         else:
-            self._unfinished = {f: i for (i, f) in enumerate(args)}
-            futures = args
+            items = args
+            self._unfinished = {
+                convert_yielded(item): i for i, item in enumerate(items)
+            }
 
         self._finished: collections.deque[Future] = collections.deque()
         self.current_index: str | int | None = None
         self.current_future: Future | None = None
         self._running_future: Future | None = None
 
-        for future in futures:
+        for future in list(self._unfinished):
             future_add_done_callback(future, self._done_callback)
 
     def done(self) -> bool:

tornado/test/gen_test.py

@@ -929,6 +929,23 @@ def test_no_ref(self):
             datetime.timedelta(seconds=0.1), gen.WaitIterator(gen.sleep(0)).next()
         )
 
+    @gen_test
+    def test_iterator_accepts_coroutine(self):
+        # WaitIterator's constructor is documented to take any
+        # awaitable, not just Futures. Passing coroutines (e.g. the
+        # result of ``gen.sleep`` or ``tornado.locks.Event().wait()``)
+        # must work without raising and produce the awaited values.
+        async def produce(value: int, delay: float) -> int:
+            await gen.sleep(delay)
+            return value
+
+        g = gen.WaitIterator(produce(1, 0.01), produce(2, 0.02), produce(3, 0.0))
+        results: dict[int, int] = {}
+        while not g.done():
+            r = yield g.next()
+            results[g.current_index] = r
+        self.assertEqual(results, {0: 1, 1: 2, 2: 3})
+
 
 class RunnerGCTest(AsyncTestCase):
     def is_pypy3(self):