1# MicroPython uasyncio module
2# MIT license; Copyright (c) 2019-2020 Damien P. George
3
4from . import core
5
6
7async def wait_for(aw, timeout, sleep=core.sleep):
8    aw = core._promote_to_task(aw)
9    if timeout is None:
10        return await aw
11
12    def runner(waiter, aw):
13        nonlocal status, result
14        try:
15            result = await aw
16            s = True
17        except BaseException as er:
18            s = er
19        if status is None:
20            # The waiter is still waiting, set status for it and cancel it.
21            status = s
22            waiter.cancel()
23
24    # Run aw in a separate runner task that manages its exceptions.
25    status = None
26    result = None
27    runner_task = core.create_task(runner(core.cur_task, aw))
28
29    try:
30        # Wait for the timeout to elapse.
31        await sleep(timeout)
32    except core.CancelledError as er:
33        if status is True:
34            # aw completed successfully and cancelled the sleep, so return aw's result.
35            return result
36        elif status is None:
37            # This wait_for was cancelled externally, so cancel aw and re-raise.
38            status = True
39            runner_task.cancel()
40            raise er
41        else:
42            # aw raised an exception, propagate it out to the caller.
43            raise status
44
45    # The sleep finished before aw, so cancel aw and raise TimeoutError.
46    status = True
47    runner_task.cancel()
48    await runner_task
49    raise core.TimeoutError
50
51
52def wait_for_ms(aw, timeout):
53    return wait_for(aw, timeout, core.sleep_ms)
54
55
56async def gather(*aws, return_exceptions=False):
57    ts = [core._promote_to_task(aw) for aw in aws]
58    for i in range(len(ts)):
59        try:
60            # TODO handle cancel of gather itself
61            # if ts[i].coro:
62            #    iter(ts[i]).waiting.push_head(cur_task)
63            #    try:
64            #        yield
65            #    except CancelledError as er:
66            #        # cancel all waiting tasks
67            #        raise er
68            ts[i] = await ts[i]
69        except Exception as er:
70            if return_exceptions:
71                ts[i] = er
72            else:
73                raise er
74    return ts
75