Проверьте, ожидали ли сопрограммы или нет - PullRequest
3 голосов
/ 25 октября 2019

У меня есть асинхронная функция, которая подключается к базе данных. В настоящее время мои пользователи делают:

conn = await connect(uri, other_params)

Я хочу продолжать поддерживать это, но хочу дополнительно разрешить использование connect() в качестве диспетчера контекста:

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

Разница междуЭти два сценария состоят в том, что в первом случае ожидается connect, а во втором - нет.

Можно ли сказать, в теле connect,ожидали ли сопрограммы или нет?

Мои текущие усилия на этом repl.it .

Ответы [ 2 ]

2 голосов
/ 26 октября 2019

Вот код, который проходит тесты, которые вы предоставили:

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

Обратите внимание, как три магических метода позволяют нам сделать объект доступным и асинхронный менеджер контекста одновременно.

Не стесняйтесь спрашиватьвопросы, если у вас есть.

0 голосов
/ 25 октября 2019

Я думаю, что в вашем примере есть структурная проблема.

Итак, сначала вашему примеру нужно 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.

...