Test if coroutine was awaited or not

user_2071807 0
Viewed 512

I have an asynchronous function which connects to a database. Currently my users do:

conn = await connect(uri, other_params)

I want to continue to support this, but want to additionally allow connect() to be used as a context manager:

async with connect(uri, other_params) as conn:
     pass

The difference between these two scenarios is that in the first case connect is awaited, and in the second case it is not.

Is it possible to tell, within the body of connect, if the coroutine was awaited or not?

My current effort at this on repl.it.

python-3.7 python-asyncio python-3.x python
2 Answers

Here's code that passes tests you provided:

import asyncio
import pytest
from functools import wraps


def connection_context_manager(func):
  @wraps(func)
  def wrapper(*args, **kwargs):

    class Wrapper:
        def __init__(self):
          self._conn = None

        async def __aenter__(self):
            self._conn = await func(*args, **kwargs)
            return self._conn

        async def __aexit__(self, *_):
          await self._conn.close()

        def __await__(self):
            return func(*args, **kwargs).__await__()  # https://stackoverflow.com/a/33420721/1113207
    return Wrapper()

  return wrapper

Note how three magic methods allows us to make object awaitable and async context manager at the same time.

Feel free to ask questions if you have any.

user_11132070

I think there is a structural issue in your example.

So first your example needs to await the __call__:

@pytest.mark.asyncio
async def test_connect_with_context_manager():
    async with await connect("with context uri") as no_context:
        # Now it's open
        assert no_context.open

    # Now it's closed
    assert not no_context.open

But here the problem is, the result of await connect("with context uri") is a Connection, which doesn't even have __aexit__ method.

So I believe you should change the structure totally, adding connect method to the Connection to actually make a connection, and in MagicConnection, delegating every method of the Connection.

user_55882790
Related