[FIXED] How to cancel all tasks in a TaskGroup

Issue

import asyncio
import random

task_group: asyncio.TaskGroup | None = None

async def coro1():
    while True:
        await asyncio.sleep(1)
        print("coro1")

async def coro2():
    while True:
        await asyncio.sleep(1)
        if random.random() < 0.1:
            print("dead")
            assert task_group is not None
            task_group.cancel() # This function does not exist.
        else:
            print("Survived another second")

async def main():
    global task_group
    async with asyncio.TaskGroup() as tg:
        task_group = tg
        tg.create_task(coro1())
        tg.create_task(coro2())
    task_group = None

asyncio.run(main())

In this example, coro1 will print "coro1" every second, coro2 has a 10% chance to cancel the entire TaskGroup, i.e., cancel both coro1 and coro2 and exit the async with block every second.

The problem is I don't know how to cancel the task group. There is no TaskGroup.cancel() function.

Solution

I end up use another task to wrap the function that contains the TaskGroup and cancel that task instead. It works as desired.

# Only show functions that has been changed

main_task: asyncio.Task[None] | None = None

async def coro2():
    while True:
        await asyncio.sleep(1)
        if random.random() < 0.1:
            print("dead")
            assert main_task is not None
            main_task.cancel()
            main_task = None
        else:
            print("Survived another second")

async def main2():
    global main_task
    main_task = asyncio.create_task(main())

asyncio.run(main2())

Answered By - Jeffrey Chen