Generator-aware decorator factory

context_entering_decorator.py

def context_entering_decorator(context_manager):
    """
    Creates a context-entering decorator that does not fail to do its job
    if the underlying function is actually a generator function.
    """

    def decorator(func):
        if inspect.isgeneratorfunction(func):
            @functools.wraps(func)
            def wrapped(*args, **kwargs):
                with context_manager():
                    yield from func(*args, **kwargs)
        else:
            @functools.wraps(func)
            def wrapped(*args, **kwargs):
                with context_manager():
                    return func(*args, **kwargs)
        return wrapped

    return decorator