Отмена только основной задачи в программе asyncio

Обычно, если сопрограмма запускается с использованием функции 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)), потому что я хотел бы, чтобы основная задача имела полный контроль над порядком отмены / выполнения своих подзадач.


python python-3.9 python-asyncio
person Jaanus Varus    schedule 15.03.2021    source источник

Ответы (1)


arrow_upward
1
arrow_downward

Не знаю, хорошая ли это идея, но в 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
person NobbyNobbs    schedule 15.03.2021