如何将contextmanager和asynccontextmanager合并到一个装饰器中？

Question

在写问题时自己解决了这个问题。橡皮鸭再次获胜

我想装饰一个生成器，既可以用作上下文管理器，也可以用作异步上下文管理器。现在我正在用一个包裹另一个

from contextlib import contextmanager, asynccontextmanager

@contextmanager
def sync_manager():
    yield

@asynccontextmanager
async def async_manager():
    with sync_manager():
        yield

但这意味着调用者需要指定他们要的是同步版本还是异步版本。我想避免这样。现在

• contextmanager的工作原理是使用__enter__和__exit__方法将函数包装在_GeneratorContextManager对象中，而
• asynccontextmanager的工作原理是使用__aenter__和__aexit__方法将函数包装在_AsyncGeneratorContextManager对象中

因为这两个协议是不同的，所以我应该能够编写一个实现这两个协议的decorator。我在这方面做了几次尝试，但运气不好。正确的方法是什么

Answer 1

解决方案：

from contextlib import contextmanager, asynccontextmanager
from functools import wraps

class MaybeAsyncGeneratorContextManager:

    def __init__(self, func, args, kwargs):
        self._func = func
        self._args = args
        self._kwargs = kwargs
        self._sync = None
        self._async = None

    def __enter__(self):
        if self._sync is None:
            syncfunc = contextmanager(self._func)
            self._sync = syncfunc(*self._args, **self._kwargs)
        return type(self._sync).__enter__(self._sync)

    def __exit__(self, t, v, tb):
        return type(self._sync).__exit__(self._sync, t, v, tb)

    def __aenter__(self):
        if self._async is None:
            @asynccontextmanager
            async def asyncfunc(*args, **kwargs):
                with contextmanager(self._func)(*args, **kwargs):
                    yield 
            self._async = asyncfunc(*self._args, **self._kwargs)
        return type(self._async).__aenter__(self._async)

    def __aexit__(self, t, v, tb):
        return type(self._async).__aexit__(self._async, t, v, tb)

def maybeasynccontextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return MaybeAsyncGeneratorContextManager(func, args, kwds)
    return helper

我最初直接使用内部的_GeneratorContextManager和_AsyncGeneratorContextManager类，但是在错误上获得正确的行为是很棘手的。这种方式添加了另一层间接寻址，但可以正确地处理错误