Await sub task finished on outer task finished event (async version of add_done_callback)

Question

Let's say we have some task (sub task) that should be finished when outer task done. We have no control of outer task: we don't know when it would be finished (it can happen before sub task done), we can't wait for sub task inside.

In this snippet we will get warning because outer task finished before sub task:

import asyncio


def create_sub_task():
    sub_task = asyncio.ensure_future(sub())
    # We want this sub_task to be finished when outer task done


async def sub():
    await asyncio.sleep(2)
    print('sub done')


async def main():  # main is outer task for sub_task
    create_sub_task()
    await asyncio.sleep(1)
    print('outer done')


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

add_done_callback looks like a way to catch moment when outer task done, but we can't wait for sub task here: this function is synchronous.

Way I found is to use event loop's private _run_once function to wait task finished inside callback synchronously:

import asyncio
from functools import partial


def create_sub_task():
    sub_task = asyncio.ensure_future(sub())

    # Callback to wait for sub_task
    outer_task = asyncio.Task.current_task()
    outer_task.add_done_callback(partial(_stop_task, sub_task))


async def sub():
    await asyncio.sleep(2)
    print('sub done')


def _stop_task(sub_task, task):
    # Ugly way to wait sub_task finished:
    loop = asyncio.get_event_loop()
    while not sub_task.done():
        loop._run_once()


async def main():  # main is outer task for sub_task
    create_sub_task()
    await asyncio.sleep(1)
    print('outer done')


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

It works, but it's ugly way with many possible problems.

Any ideas how to solve task better?

Answer 1

AFAIK there is no way to solve this without internals. Personally I would asyncio.gather outer and sub tasks in one future and then rewrite callbacks.

Unfortunately Future's list of callbacks is not exposed with public interface (I'm using _callbacks):

import asyncio

def create_sub_task():
    sub_task = asyncio.ensure_future(sub())
    outer_task = asyncio.Task.current_task()

    multi_fut = asyncio.gather(sub_task, outer_task)
    for cb in outer_task._callbacks:
        multi_fut.add_done_callback(cb)
        outer_task.remove_done_callback(cb)

async def sub():
    await asyncio.sleep(2)
    print('sub done')


async def main():  # main is outer task for sub_task
    create_sub_task()
    await asyncio.sleep(1)
    print('outer done')


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

I assume that you don't want or you can't change the flow, but I encourage you to rethink. Maybe post some context - constraints origin.