From ecf3d3a275929f4fd3be25783b1e7fa5a1477eb9 Mon Sep 17 00:00:00 2001 From: Zo Bot Date: Thu, 2 Jul 2026 17:32:40 +0000 Subject: [PATCH] gen: accept any _Yieldable in WaitIterator constructor --- tornado/gen.py | 13 +++++++------ tornado/test/gen_test.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tornado/gen.py b/tornado/gen.py index 4cd3ba273..ce6453539 100644 --- a/tornado/gen.py +++ b/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: diff --git a/tornado/test/gen_test.py b/tornado/test/gen_test.py index 99167f718..7ddbdeeb4 100644 --- a/tornado/test/gen_test.py +++ b/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):