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中委派{{1}的每种方法}。

测试是否等待协程

2 个答案: