How to return a value from a cancelled python asyncio coroutine due to timeout

Question

In python > 3.5, how can I have a coroutine return a final value after being cancelled due to TimeoutError?

I have a small python project that uses multiple coroutines to transfer data and reports back the amount of data transferred. It takes a timeout parameter; if the script times out prior to completing the transfer, it reports back the amount that it transferred prior to cancellation.

It was working fine in python3.5, but recently I tried updating to 3.8 and ran into trouble.

Below is example code, and clearly its behavior differs widely from 3.5, 3.6, 3.7, and 3.8:

import asyncio
import sys


async def foo():
    try:
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print("foo got cancelled")
        return 1


async def main():
    coros = asyncio.gather(*(foo() for _ in range(3)))
    try:
        await asyncio.wait_for(coros, timeout=0.1)
    except asyncio.TimeoutError:
        print("main coroutine timed out")
        await coros
    return coros.result()


if __name__ == "__main__":
    print(sys.version)

    loop = asyncio.new_event_loop()
    try:
        results = loop.run_until_complete(main())
        print("results: {}".format(results))
    except Exception as e:
        print("exception in __main__:")
        print(e)
    finally:
        loop.close()

$ for ver in 3.5 3.6 3.7 3.8; do echo; python${ver} example.py; done

3.5.7 (default, Sep  6 2019, 07:49:56)
[GCC 4.2.1 Compatible Apple LLVM 10.0.1 (clang-1001.0.46.4)]
main coroutine timed out
foo got cancelled
foo got cancelled
foo got cancelled
results: [1, 1, 1]

3.6.9 (default, Sep  6 2019, 07:45:14)
[GCC 4.2.1 Compatible Apple LLVM 10.0.1 (clang-1001.0.46.4)]
main coroutine timed out
foo got cancelled
foo got cancelled
foo got cancelled
exception in __main__:


3.7.4 (default, Sep 17 2019, 13:46:30)
[Clang 10.0.1 (clang-1001.0.46.4)]
foo got cancelled
foo got cancelled
foo got cancelled
main coroutine timed out
exception in __main__:


3.8.0 (default, Oct 16 2019, 21:30:17)
[Clang 11.0.0 (clang-1100.0.33.8)]
foo got cancelled
foo got cancelled
foo got cancelled
main coroutine timed out
Traceback (most recent call last):
  File "example.py", line 28, in <module>
    results = loop.run_until_complete(main())
  File "/usr/local/var/pyenv/versions/3.8.0/lib/python3.8/asyncio/base_events.py", line 608, in run_until_complete
    return future.result()
asyncio.exceptions.CancelledError

exception in __main__: is not printing for 3.8 because CancelledError is now BaseException instead of Exception (EDIT: which may be why the traceback prints here but not elsewhere).

I have tried a number of configurations of using return_exceptions=True in asyncio.gather or catching CancelledError in the except asyncio.TimeoutError: block, but I can't seem to get it right.

I need to keep main as an async function, because in my actual code it is creating an aiohttp Session for the other coroutines to share, and modern aiohttp requires this to be done in an async contextmanager (instead of a regular sync context manager).

I'm hoping for code that runs on 3.5-3.8, so I'm not using asyncio.run.

I've tried code from a number of other questions that use .cancel() with or without contextlib.suppress(asyncio.CancelledError), but still no luck. I've also tried returning an awaited value (e.g. result = await coros; return result instead of return coros.result()), also no dice.

Is there a good way for me to get the python 3.5 behavior in python >3.5, in which I can have a coroutine catch CancelledError on timeout and return a value when next awaited?

Thanks in advance.

Answer 1

Thanks to @RafalS and their suggestion to stop using asyncio.gather.

Instead of using gather and wait_for, it seems that using the timeout from .wait directly with the coroutines may be the best bet, and works from 3.5 to 3.8.

Note that the bash command below is slightly modified to show that the tasks are being run simultaneously and also being cancelled without waiting for foo to complete.

import asyncio
import sys


async def foo():
    try:
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        pass
    finally:
        return 1


async def main():
    coros = [foo() for _ in range(3)]
    done, pending = await asyncio.wait(coros, timeout=1.0)
    for task in pending:
        task.cancel()
        await task
    return [task.result() for task in done | pending]


if __name__ == "__main__":
    print(sys.version)

    loop = asyncio.new_event_loop()
    try:
        results = loop.run_until_complete(main())
        print("results: {}".format(results))
    finally:
        loop.close()

$ for ver in 3.5 3.6 3.7 3.8; do echo; time python${ver} example.py; done

3.5.7 (default, Sep  6 2019, 07:49:56)
[GCC 4.2.1 Compatible Apple LLVM 10.0.1 (clang-1001.0.46.4)]
results: [1, 1, 1]

real    0m1.634s
user    0m0.173s
sys     0m0.106s

3.6.9 (default, Sep  6 2019, 07:45:14)
[GCC 4.2.1 Compatible Apple LLVM 10.0.1 (clang-1001.0.46.4)]
results: [1, 1, 1]

real    0m1.643s
user    0m0.184s
sys     0m0.100s

3.7.4 (default, Sep 17 2019, 13:46:30)
[Clang 10.0.1 (clang-1001.0.46.4)]
results: [1, 1, 1]

real    0m1.499s
user    0m0.129s
sys     0m0.089s

3.8.0 (default, Oct 16 2019, 21:30:17)
[Clang 11.0.0 (clang-1100.0.33.8)]
results: [1, 1, 1]

real    0m1.492s
user    0m0.141s
sys     0m0.087s