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

Question

通常，如果使用asyncio.run(coroutine) function 启动协程，则键盘中断 (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) ），因为我希望主要任务能够完全控制其子任务的取消/执行顺序。

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