bpo-40816 Add AsyncContextDecorator class (GH-20516)

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

1 files changed, 24 insertions, 1 deletions

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__()