How to process tasks as they complete when using TaskGroup?

[odigity]

I understand the arguments for using the newer TaskGroup in place of older mechanisms based on create_task().

However, TaskGroup exits after all tasks are done. What if you want to start processing the results as soon as the first task finishes, like you can with asyncio.as_completed() or asyncio.wait( ..., return_when=FIRST_COMPLETED )?

If TaskGroup offers no way to do that, then isn't using it a trade-off rather than a strict upgrade?

[DeaD_EyE]

It's possible to replicate similar behavior.
Example:

import asyncio


class TaskException(Exception):
    pass


async def worker(n):
    try:
        await asyncio.sleep(n)
    except asyncio.CancelledError:
        print(f"Task {n} has been cancelled")
        # Task 10 has been cancelled

    if n == 9:
        raise TaskException(f"Exception in Task {n}")

    return n


async def main():
    tasks_working = set()
    tasks_done = set()

    def task_done(task):
        tasks_working.discard(task)
        tasks_done.add(task)
        

    async with asyncio.TaskGroup() as tg:
        for n in range(1,11):
            task = tg.create_task(worker(n))
            tasks_working.add(task)
            task.add_done_callback(task_done)

        while tasks_working or tasks_done:
            print(".", end="", flush=True)
            
            if tasks_done:
                print()
                
            for task in tasks_done.copy():
                print(f"Task {task.result()} is done")
                tasks_done.discard(task)
                
            await asyncio.sleep(1)
            
        print("All tasks are done")
    print("Leaving TGs context")

try:
    asyncio.run(main())
except *TaskException as e:
    for exc in e.exceptions:
        print(exc)
        # Exception in Task 9

Output:

Output:
...
Task 2 is done
Task 1 is done
.
Task 3 is done
.
Task 4 is done
.
Task 5 is done
.
Task 6 is done
.
Task 7 is done
.
Task 8 is done
Task 10 has been cancelled
Exception in Task 9

Hm, print("All tasks are done") and print("Leaving TGs context") are not called.
I thought if there is an Exception inside the TaskGroup, all Tasks are cancelled.
I think it's not called, because the Exceptions happens in the eventloop somewhere and ends it before the eventloop is able to execute the rest.

[odigity]

I didn't quite follow of all that, but the gist of it seems to be "use task.add_done_callback() inside the TG context to implement timely task result processing without waiting for TG context to exit".

Thanks, I'll give that pattern a try.