如何将contextmanager和异步contextmanager合并为一个装饰器？

Question

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

我想装饰一个生成器，作为上下文管理器和异步上下文管理器。现在我把一个包在另一个里

from contextlib import contextmanager, asynccontextmanager

@contextmanager
def sync_manager():
    yield

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

但这意味着调用方需要指定它们所追求的同步或异步版本。我想避免那样做。现在

• contextmanager通过用__enter__和__exit__方法包装_GeneratorContextManager对象中的函数来工作，while
• asynccontextmanager通过用__aenter__和__aexit__方法包装_AsyncGeneratorContextManager对象来工作。

由于它们是不同的，所以我应该能够编写一个实现这两种协议的装饰器。我在这方面做了几次没有运气的尝试。正确的方法是什么？

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类，但在错误上获得正确的行为是很棘手的。这种方式增加了另一层间接，但正确地处理错误。