bpo-29679: Implement @contextlib.asynccontextmanager (#360)

Author: JelleZijlstra

Doc/library/contextlib.rst

@@ -80,6 +80,36 @@ Functions and classes provided:
       Use of :class:`ContextDecorator`.
 
 
+.. decorator:: asynccontextmanager
+
+   Similar to :func:`~contextlib.contextmanager`, but creates an
+   :ref:`asynchronous context manager <async-context-managers>`.
+
+   This function is a :term:`decorator` that can be used to define a factory
+   function for :keyword:`async with` statement asynchronous context managers,
+   without needing to create a class or separate :meth:`__aenter__` and
+   :meth:`__aexit__` methods. It must be applied to an :term:`asynchronous
+   generator` function.
+
+   A simple example::
+
+      from contextlib import asynccontextmanager
+
+      @asynccontextmanager
+      async def get_connection():
+          conn = await acquire_db_connection()
+          try:
+              yield
+          finally:
+              await release_db_connection(conn)
+
+      async def get_all_users():
+          async with get_connection() as conn:
+              return conn.query('SELECT ...')
+
+   .. versionadded:: 3.7
+
+
 .. function:: closing(thing)
 
    Return a context manager that closes *thing* upon completion of the block.  This

Doc/reference/datamodel.rst

@@ -2575,6 +2575,8 @@ An example of an asynchronous iterable object::
       result in a :exc:`RuntimeError`.
 
 
+.. _async-context-managers:
+
 Asynchronous Context Managers
 -----------------------------
 

Doc/whatsnew/3.7.rst

@@ -95,6 +95,12 @@ New Modules
 Improved Modules
 ================
 
+contextlib
+----------
+
+:func:`contextlib.asynccontextmanager` has been added. (Contributed by
+Jelle Zijlstra in :issue:`29679`.)
+
 distutils
 ---------
 

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 @@ 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.