检验是否等待协同疗法

Question

我有一个连接到数据库的异步函数。目前，我的用户做：

conn = await connect(uri, other_params)

我希望继续支持这一点，但希望另外允许将connect()用作上下文管理器：

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

这两种场景的不同之处在于，在第一种情况下，等待的是connect，而在第二种情况下，则不是。

在connect的身体内，是否有可能知道是否等待着协同线？

我目前在repl.it上所做的努力。

Answer 1

下面是通过您提供的测试的代码：

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

请注意，三种神奇的方法允许我们同时使对象成为可接受的和异步的上下文管理器。

如果你有问题的话，可以随意提问。

Answer 2

我认为你的例子有一个结构性的问题。

因此，首先，您的示例需要await __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

但是这里的问题是，await connect("with context uri")的结果是一个Connection，它甚至没有__aexit__方法。

因此，我认为您应该完全改变结构，将connect方法添加到Connection中以实际建立连接，并在MagicConnection中委托Connection的每个方法。