仅取消异步程序中的主要任务

Question

通常情况下，如果使用asyncio.run(coroutine)函数启动协程，键盘中断（CTRL + C）或SIGINT将取消事件循环中所有待处理的任务。我正在寻找一种只取消主要任务（传递给 asyncio.run(coroutine) 的任务）的方法。这个想法是，主任务将按照它认为合适的顺序安排所有子任务的取消。

考虑一个例子：

import asyncio


async def main():
    foo_task = asyncio.create_task(foo())
    try:
        await asyncio.sleep(10)
        print('main finished')
    finally:
        print('ensuring foo task is finished')
        await foo_task


async def foo():
    await asyncio.sleep(10)
    print('foo finished')


try:
    asyncio.run(main())
except KeyboardInterrupt:
    pass

我想更改上面的代码，以便如果在执行过程中发送键盘中断或 SIGINT，foo_task 仍会完成。它应该打印以下内容：

ensuring foo task is finished
foo finished

我不想使用屏蔽(asyncio.shield(coroutine))，因为我希望主任务完全控制其子任务cancellation/execution的顺序。

Answer 1

我不知道这是否是个好主意，但在类 Unix 操作系统上，您可以使用信号处理程序实现所需的行为。

import asyncio
from asyncio import tasks
import signal
from typing import Coroutine, Set

to_cancel: Set[Coroutine] = set()  # little workaround to detect the main task

async def main():
    loop = asyncio.get_event_loop()
    loop.add_signal_handler(signal.SIGINT, cancel_main)
    loop.add_signal_handler(signal.SIGTERM, cancel_main)

    foo_task = asyncio.create_task(foo())

    try:
        print("main sleeping")
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print("main cancelled")
    finally:
        print('ensuring foo task is finished')
        await foo_task
        print('main finished')


async def foo():
    print("foo sleeping")
    await asyncio.sleep(10)
    print("foo finished")


def cancel_main():
    for task in tasks.all_tasks():
        # task.get_coro() for python >= 3.8 else task._coro
        if task.get_coro() in to_cancel and not task.cancelled():
            task.cancel()

if __name__ == "__main__":
    coro = main()
    to_cancel.add(coro)
    asyncio.run(coro)

结果

main sleeping
foo sleeping
^C
main cancelled
ensuring foo task is finished
foo finished
main finished

仅取消异步程序中的主要任务

Cancelling only the main task in an asyncio program

python

python-asyncio

python-3.9