intercepting the call to __await__ on an asyncio Task in Python

Question

What I'm trying to do is to get notified whenever:

an asyncio Task is created
anyone awaits on an asyncio Task

The first one is straightforward - use a task factory. For the second one - there are two implementation approaches I could think of:

1) Wrap the task with a new Python object that behaves like a future (enough to fool other code), forward whatever methods necessary to the wrapped task

2) monkey patch await in the original task to notify me, and then forward it to the original __await.

I tried to go for option (2) first because it seems easier. However, I can't get it to work. Attempts to modify await (directly or via setattr) seem to have no effect whatsoever.

Questions:

Any way to get (2) to work?
How safe is the (1) approach?
Is there a "clean" way to wrap a task in a future that I could use?

Here's some code showing what I tried doing.

import asyncio

def patch_task_factory(loop):
    prior = loop.get_task_factory()
    def factory(loop, coro):
        if prior is None:
            task = asyncio.Task(coro, loop=loop)
        else:
            task = prior(loop, coro)

        # here we know a task has been created
        # but to also get notified when someone waits for it, we need to patch __await__
        orig_await = task.__await__

        def intercepted_await(self):
            print("someone is awaiting this task")
            # should work because we're just returning the iterator from the real await
            return orig_await(self)
        task.__await__ = intercepted_await.__get__(task, asyncio.Task)

        return task
    loop.set_task_factory(factory)


async def test1():
    print("test1!")
    return 1

async def test2():
    print("test2!")
    return await asyncio.create_task(test1())


def main():
    loop = asyncio.get_event_loop()
    patch_task_factory(loop)
    res = asyncio.run(test2())
    print(res)


if __name__ == "__main__":
    # execute only if run as a script
    main()

user_1523042 · Answer

The issue is that Task is a built-in/extension method, imported from _asyncio (note the underscore). You'll note that setattr(task.__class__.__getattribute__, ...), for example, complains that you cannot modify method-wrappers of built-in classes.

You might consider a proxy object in this case. Note the "trickery" around overriding __await__:

proxy.py is taken straight from http://code.activestate.com/recipes/496741-object-proxying/.

$ cat task_factory.py

#!/usr/bin/env python3

import asyncio
from typing import Any, Generator
from proxy import Proxy


class ProxyWithAwait(Proxy):
    def __await__(self) -> Generator[Any, None, Any]:
        _obj = object.__getattribute__(self, "_obj")
        aw = _obj.__await__()
        print("*** INTERCEPTING AWAIT: BEFORE")
        yield aw.send(None)
        print("*** INTERCEPTING AWAIT: AFTER")
        aw.close()


async def main() -> None:

    loop = asyncio.get_event_loop()
    old_task_factory = loop.get_task_factory()

    def my_task_factory(
        loop: asyncio.AbstractEventLoop,
        coro: Generator[Any, None, Any],
    ) -> asyncio.Future:
        if old_task_factory is None:
            t = asyncio.tasks.Task(coro, loop=loop)
        else:
            t = old_task_factory(loop, coro)

        return ProxyWithAwait(t)

    loop.set_task_factory(my_task_factory)

    async def my_task_func() -> None:
        print("my_task_func(): start")
        await asyncio.sleep(1.5)
        print("my_task_func(): end")

    t = loop.create_task(my_task_func())
    print("main: before await t")
    await t
    print("main: after await t")

    loop.set_task_factory(old_task_factory)


if __name__ == "__main__":
    asyncio.run(main())

intercepting the call to await on an asyncio Task in Python

1 Answers