Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

bpo-40816 Add AsyncContextDecorator class #20516

Merged
merged 12 commits into from
Nov 5, 2020
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion Lib/contextlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,22 @@ def inner(*args, **kwds):
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."""

Expand Down Expand Up @@ -167,9 +183,16 @@ def __exit__(self, type, value, traceback):


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__()
Expand Down
27 changes: 27 additions & 0 deletions Lib/test/test_contextlib_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -278,6 +278,33 @@ async def woohoo(self, func, args, kwds):
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):
heckad marked this conversation as resolved.
Show resolved Hide resolved
depth = 0
heckad marked this conversation as resolved.
Show resolved Hide resolved
ncols = 0

@asynccontextmanager
async def woohoo():
nonlocal ncols
ncols += 1

nonlocal depth
before = depth
depth += 1
heckad marked this conversation as resolved.
Show resolved Hide resolved
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)
heckad marked this conversation as resolved.
Show resolved Hide resolved


class TestAsyncExitStack(TestBaseExitStack, unittest.TestCase):
class SyncAsyncExitStack(AsyncExitStack):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add AsyncContextDecorator to contextlib to support async context manager as a decorator.