bpo-29679: Implement @contextlib.asynccontextmanager (#360)
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index 5e47054..c53b35e 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -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 @@
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 @@
# 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 @@
# 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,7 +208,6 @@
<body>
finally:
<cleanup>
-
"""
@wraps(func)
def helper(*args, **kwds):
@@ -161,6 +215,39 @@
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.