Skip to content

Instantly share code, notes, and snippets.

@belm0
Last active July 25, 2019 06:58
Show Gist options
  • Save belm0/a0e3c6ea6402b177ca22dca9ca29ad89 to your computer and use it in GitHub Desktop.
Save belm0/a0e3c6ea6402b177ca22dca9ca29ad89 to your computer and use it in GitHub Desktop.
Python decorator for creating *non-async* context managers which can decorate *async* functions
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment