Show task stack for current task, and coroutine stack for each task therein

tstack.py

"""Reconstruct awaiters stack, including Tasks."""

import asyncio
import os


def get_await_stack(coro):
    """Return a coroutine's chain of awaiters.
    
    This follows the cr_await links.
    """
    stack = []
    while coro is not None and hasattr(coro, "cr_await"):
        stack.append(coro)
        coro = coro.cr_await
    return stack


def get_task_tree():
    """Return the task tree dict {awaited: awaiting}.

    This follows the _fut_waiter links and constructs a map
    from awaited tasks to the tasks that await them.
    """
    tree = {}
    for task in asyncio.all_tasks():
        awaited = task._fut_waiter
        if awaited is not None:
            tree[awaited] = task
    return tree


def get_task_stack(task):
    """Return the stack of tasks awaiting a task.

    For each task it returns a tuple (task, awaiters) where
    awaiters is the chain of coroutines comprising the task.
    The first entry is the argument task, the last entry is
    the root task (often "Task-1", created by asyncio.run()).

    For example, if we have a task A running a coroutine f1,
    where f1 awaits f2, and f2 awaits a task B running a coroutine
    f3 which awaits f4 which awaits f5, then we'll return

        [
            (B, [f3, f4, f5]),
            (A, [f1, f2]),
        ]

    NOTE: The coroutine stack for the *current* task is inaccessible.
    To work around this, use `await task_stack()`.

    TODO:
    - Maybe it would be nicer to reverse the stack?
    - Classic coroutines and async generators are not supported yet.
    - This is expensive due to the need to first create a reverse
      mapping of awaited tasks to awaiting tasks.
    """
    tree = get_task_tree()
    stack = []
    while task is not None:
        coro = task.get_coro()
        awaiters = get_await_stack(coro)
        stack.append((task, awaiters))
        task = tree.get(task)
    return stack


async def task_stack():
    """Return the stack of tasks awaiting the current task.

    This exists so you can get the coroutine stack for the current task.
    """
    task = asyncio.current_task()
    async def helper(task):
        return get_task_stack(task)
    return await asyncio.create_task(helper(task), name="TaskStackHelper")


# Example code

async def f1():
    await f2()

async def f2():
    await asyncio.create_task(f3(), name="F3")

async def f3():
    await asyncio.create_task(f4(), name="F4")

async def f4():
    await f5()

async def f5():
    tasks = await task_stack()
    # tasks = get_task_stack(asyncio.current_task())
    report_task_stack(tasks)


def report_task_stack(tasks):
    """Helper to summarize a task stack."""
    for task, awaiters in reversed(tasks):
        print(f"Task: {task.get_name()}")
        for awaiter in awaiters:
            print(f"  Coro: {awaiter.__qualname__} {os.path.relpath(awaiter.cr_code.co_filename)}:{awaiter.cr_code.co_firstlineno}")


async def main():
    await asyncio.create_task(f1(), name="F1")


asyncio.run(main())