bpo-40816 Add AsyncContextDecorator class (GH-20516)

Co-authored-by: Yury Selivanov <yury@edgedb.com>

4 files changed, 114 insertions, 1 deletions

diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
index e42f5a9..ee2becb 100644
--- a/Doc/library/contextlib.rst
+++ b/Doc/library/contextlib.rst
@@ -126,6 +126,31 @@ Functions and classes provided:
 
    .. versionadded:: 3.7
 
+   Context managers defined with :func:`asynccontextmanager` can be used
+   either as decorators or with :keyword:`async with` statements::
+
+     import time
+
+     async def timeit():
+         now = time.monotonic()
+         try:
+             yield
+         finally:
+             print(f'it took {time.monotonic() - now}s to run')
+
+      @timeit()
+      async def main():
+          # ... async code ...
+
+   When used as a decorator, a new generator instance is implicitly created on
+   each function call. This allows the otherwise "one-shot" context managers
+   created by :func:`asynccontextmanager` to meet the requirement that context
+   managers support multiple invocations in order to be used as decorators.
+
+  .. versionchanged:: 3.10
+     Async context managers created with :func:`asynccontextmanager` can
+     be used as decorators.
+
 
 .. function:: closing(thing)
 
@@ -384,6 +409,43 @@ Functions and classes provided:
    .. versionadded:: 3.2
 
 
+.. class:: AsyncContextManager
+
+   Similar as ContextManger only for async
+
+   Example of ``ContextDecorator``::
+
+      from asyncio import run
+      from contextlib import AsyncContextDecorator
+
+      class mycontext(AsyncContextDecorator):
+          async def __aenter__(self):
+              print('Starting')
+              return self
+
+          async def __aexit__(self, *exc):
+              print('Finishing')
+              return False
+
+      >>> @mycontext()
+      ... async def function():
+      ...     print('The bit in the middle')
+      ...
+      >>> run(function())
+      Starting
+      The bit in the middle
+      Finishing
+
+      >>> async def function():
+      ...    async with mycontext():
+      ...         print('The bit in the middle')
+      ...
+      >>> run(function())
+      Starting
+      The bit in the middle
+      Finishing
+
+
 .. class:: ExitStack()
 
    A context manager that is designed to make it easy to programmatically
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index 82ddc14..56b4968 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -80,6 +80,22 @@ class ContextDecorator(object):
         return inner
 
 
+class AsyncContextDecorator(object):
+    "A base class or mixin that enables async context managers to work as decorators."
+
+    def _recreate_cm(self):
+        """Return a recreated instance of self.
+        """
+        return self
+
+    def __call__(self, func):
+        @wraps(func)
+        async def inner(*args, **kwds):
+            async with self._recreate_cm():
+                return await func(*args, **kwds)
+        return inner
+
+
 class _GeneratorContextManagerBase:
     """Shared functionality for @contextmanager and @asynccontextmanager."""
 
@@ -167,9 +183,16 @@ class _GeneratorContextManager(_GeneratorContextManagerBase,
 
 
 class _AsyncGeneratorContextManager(_GeneratorContextManagerBase,
-                                    AbstractAsyncContextManager):
+                                    AbstractAsyncContextManager,
+                                    AsyncContextDecorator):
     """Helper for @asynccontextmanager."""
 
+    def _recreate_cm(self):
+        # _AGCM instances are one-shot context managers, so the
+        # ACM must be recreated each time a decorated function is
+        # called
+        return self.__class__(self.func, self.args, self.kwds)
+
     async def __aenter__(self):
         try:
             return await self.gen.__anext__()
diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py
index 3765f6c..109807d 100644
--- a/Lib/test/test_contextlib_async.py
+++ b/Lib/test/test_contextlib_async.py
@@ -278,6 +278,33 @@ class AsyncContextManagerTestCase(unittest.TestCase):
         async with woohoo(self=11, func=22, args=33, kwds=44) as target:
             self.assertEqual(target, (11, 22, 33, 44))
 
+    @_async_test
+    async def test_recursive(self):
+        depth = 0
+        ncols = 0
+
+        @asynccontextmanager
+        async def woohoo():
+            nonlocal ncols
+            ncols += 1
+
+            nonlocal depth
+            before = depth
+            depth += 1
+            yield
+            depth -= 1
+            self.assertEqual(depth, before)
+
+        @woohoo()
+        async def recursive():
+            if depth < 10:
+                await recursive()
+
+        await recursive()
+
+        self.assertEqual(ncols, 10)
+        self.assertEqual(depth, 0)
+
 
 class AclosingTestCase(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/Library/2020-05-29-15-25-41.bpo-40816.w61Pob.rst b/Misc/NEWS.d/next/Library/2020-05-29-15-25-41.bpo-40816.w61Pob.rst
new file mode 100644
index 0000000..66b7577
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-05-29-15-25-41.bpo-40816.w61Pob.rst
@@ -0,0 +1 @@
+Add AsyncContextDecorator to contextlib to support async context manager as a decorator.
\ No newline at end of file