Mixinig async context manager and straight await in asyncio

Question

How to mix

async with api.open() as o:
    ...

and

o = await api.open()

in one function?

Since first require object with __aenter__ and __aexit__, but second require __await__, which should be generator without await.

My attempt was:

def AsyncContext(aenter, aexit):

    class AsyncContextClass:

        async def __aenter__(self):

            return await aenter()

        async def __aexit__(self, *args):

            return await aexit(*args)

        def __await__(self):

            return (yield from aenter())

    return AsyncContextClass()

But it fails with __await__ if aenter defined with async def (TypeError: cannot 'yield from' a coroutine object in a non-coroutine generator).

It works fine with @asyncio.coroutine decorator for aenter, but this is «dirty».

Answer 1

You can return __aenter__'s __await__ from your class's __await__:

# vim: tabstop=4 expandtab

import asyncio

class Test(object):

    async def __aenter__(self):
        print("enter")

    async def __aexit__(self, *args):
        print("exit")

    def __await__(self):
        return self.__aenter__().__await__()

async def run():
    async with Test():
        print("hello")

    await Test()

loop = asyncio.get_event_loop()
loop.run_until_complete(run())
loop.close()

Output:

enter
hello
exit
enter