bpo-29679: Implement @contextlib.asynccontextmanager by JelleZijlstra · Pull Request #360 · python/cpython

python

@@ -4,9 +4,9 @@ from collections import deque from functools import wraps
__all__ = ["contextmanager", "closing", "AbstractContextManager", "ContextDecorator", "ExitStack", "redirect_stdout", "redirect_stderr", "suppress"] __all__ = ["asynccontextmanager", "contextmanager", "closing", "AbstractContextManager", "ContextDecorator", "ExitStack", "redirect_stdout", "redirect_stderr", "suppress"]

class AbstractContextManager(abc.ABC): @@ -54,8 +54,8 @@ def inner(*args, **kwds): return inner

class _GeneratorContextManager(ContextDecorator, AbstractContextManager): """Helper for @contextmanager decorator.""" class _GeneratorContextManagerBase: """Shared functionality for @contextmanager and @asynccontextmanager."""
def __init__(self, func, args, kwds): self.gen = func(*args, **kwds) @@ -71,6 +71,12 @@ def __init__(self, func, args, kwds): # for the class instead. # See http://bugs.python.org/issue19404 for more details.

class _GeneratorContextManager(_GeneratorContextManagerBase, AbstractContextManager, ContextDecorator): """Helper for @contextmanager decorator."""
def _recreate_cm(self): # _GCM instances are one-shot context managers, so the # CM must be recreated each time a decorated function is @@ -121,12 +127,61 @@ def __exit__(self, type, value, traceback): # fixes the impedance mismatch between the throw() protocol # and the __exit__() protocol. # # This cannot use 'except BaseException as exc' (as in the # async implementation) to maintain compatibility with # Python 2, where old-style class exceptions are not caught # by 'except BaseException'. if sys.exc_info()[1] is value: return False raise raise RuntimeError("generator didn't stop after throw()")

class _AsyncGeneratorContextManager(_GeneratorContextManagerBase): """Helper for @asynccontextmanager."""
async def __aenter__(self): try: return await self.gen.__anext__() except StopAsyncIteration: raise RuntimeError("generator didn't yield") from None
async def __aexit__(self, typ, value, traceback): if typ is None: try: await self.gen.__anext__() except StopAsyncIteration: return else: raise RuntimeError("generator didn't stop") else: if value is None: value = typ() # See _GeneratorContextManager.__exit__ for comments on subtleties # in this implementation try: await self.gen.athrow(typ, value, traceback) raise RuntimeError("generator didn't stop after throw()") except StopAsyncIteration as exc: return exc is not value except RuntimeError as exc: if exc is value: return False # Avoid suppressing if a StopIteration exception # was passed to throw() and later wrapped into a RuntimeError # (see PEP 479 for sync generators; async generators also # have this behavior). But do this only if the exception wrapped # by the RuntimeError is actully Stop(Async)Iteration (see # issue29692). if isinstance(value, (StopIteration, StopAsyncIteration)): if exc.__cause__ is value: return False raise except BaseException as exc: if exc is not value: raise

def contextmanager(func): """@contextmanager decorator.
@@ -153,14 +208,46 @@ def some_generator(<arguments>): <body> finally: <cleanup>
""" @wraps(func) def helper(*args, **kwds): return _GeneratorContextManager(func, args, kwds) return helper

def asynccontextmanager(func): """@asynccontextmanager decorator.
Typical usage:
@asynccontextmanager async def some_async_generator(<arguments>): <setup> try: yield <value> finally: <cleanup>
This makes this:
async with some_async_generator(<arguments>) as <variable>: <body>
equivalent to this:
<setup> try: <variable> = <value> <body> finally: <cleanup> """ @wraps(func) def helper(*args, **kwds): return _AsyncGeneratorContextManager(func, args, kwds) return helper

class closing(AbstractContextManager): """Context to automatically close something at the end of a block.