summaryrefslogtreecommitdiffstats
path: root/Lib
diff options
context:
space:
mode:
Diffstat (limited to 'Lib')
-rw-r--r--Lib/contextlib.py25
-rw-r--r--Lib/test/test_contextlib_async.py27
2 files changed, 51 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__()
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):