implement @asynccontextmanager

Needs docs and tests
python · 1st1 · May 1, 2017 · Jan 19, 2017 · Feb 28, 2017 · Feb 28, 2017
commit 904e8a8bb699b6eff1f086a2e60f36303d614968
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
@@ -54,8 +54,9 @@ def inner(*args, **kwds):
         return inner
 
 
-class _GeneratorContextManager(ContextDecorator, AbstractContextManager):
-    """Helper for @contextmanager decorator."""
+class _GeneratorContextManagerBase(ContextDecorator):
+    """Shared functionality for the @contextmanager and @asynccontextmanager
+    implementations."""
 
     def __init__(self, func, args, kwds):
         self.gen = func(*args, **kwds)
@@ -77,6 +78,10 @@ def _recreate_cm(self):
         # called
         return self.__class__(self.func, self.args, self.kwds)
 
+
+class _GeneratorContextManager(_GeneratorContextManagerBase, AbstractContextManager):
+    """Helper for @contextmanager decorator."""
+
     def __enter__(self):
         try:
             return next(self.gen)
@@ -126,6 +131,44 @@ def __exit__(self, type, value, traceback):
                     raise
 
 
+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, type, value, traceback):
+        if type is None:
+            try:
+                await self.gen.__anext__()
+            except StopAsyncIteration:
+                return
+            else:
+                raise RuntimeError("generator didn't stop")
+        else:
+            if value is None:
+                value = type()
+            # See _GeneratorContextManager.__exit__ for comments on subtleties
+            # in this implementation
+            try:
+                await self.gen.athrow(type, 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
+                if exc.__cause__ is value:
+                    return False
+                raise
+            except:
+                if sys.exc_info()[1] is not value:
+                    raise
+
+
 def contextmanager(func):
     """@contextmanager decorator.
 
@@ -160,6 +203,40 @@ def helper(*args, **kwds):
     return helper
 
 
+def asynccontextmanager(func):
+    """@contextmanager 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.