Python decorator for creating *non-async* context managers which can decorate *async* functions

async_friendly_contextmanager.py

from contextlib import _GeneratorContextManager
from functools import wraps
from inspect import iscoroutinefunction

class _AsyncFriendlyGeneratorContextManager(_GeneratorContextManager):
    """
    Hack contextlib._GeneratorContextManager so that resulting context
    manager can properly decorate async functions.
    """

    def __call__(self, func):
        if iscoroutinefunction(func):
            wraps(func)
            async def inner(*args, **kwargs):
                with self._recreate_cm():  # pylint: disable=not-context-manager
                    return await func(*args, **kwargs)
        else:
            wraps(func)
            def inner(*args, **kwargs):
                with self._recreate_cm():  # pylint: disable=not-context-manager
                    return func(*args, **kwargs)
        return inner


def async_friendly_contextmanager(func):
    """
    Equivalent to @contextmanager, except the resulting (non-async) context manager works
    correctly as a decorator on async functions.
    """
    @wraps(func)
    def helper(*args, **kwargs):
        return _AsyncFriendlyGeneratorContextManager(func, args, kwargs)
    return helper