Skip to content

Commit

Permalink
bpo-38415: Allow using @asynccontextmanager-made ctx managers as deco…
Browse files Browse the repository at this point in the history
…rators (GH-16667)
  • Loading branch information
fried authored Sep 23, 2021
1 parent af90b54 commit 86b833b
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 0 deletions.
8 changes: 8 additions & 0 deletions Lib/contextlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,14 @@ class _AsyncGeneratorContextManager(
):
"""Helper for @asynccontextmanager decorator."""

def __call__(self, func):
@wraps(func)
async def inner(*args, **kwds):
async with self.__class__(self.func, self.args, self.kwds):
return await func(*args, **kwds)

return inner

async def __aenter__(self):
# do not keep args and kwds alive unnecessarily
# they are only needed for recreation, which is not possible anymore
Expand Down
76 changes: 76 additions & 0 deletions Lib/test/test_contextlib_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,82 @@ async def recursive():
self.assertEqual(ncols, 10)
self.assertEqual(depth, 0)

@_async_test
async def test_decorator(self):
entered = False

@asynccontextmanager
async def context():
nonlocal entered
entered = True
yield
entered = False

@context()
async def test():
self.assertTrue(entered)

self.assertFalse(entered)
await test()
self.assertFalse(entered)

@_async_test
async def test_decorator_with_exception(self):
entered = False

@asynccontextmanager
async def context():
nonlocal entered
try:
entered = True
yield
finally:
entered = False

@context()
async def test():
self.assertTrue(entered)
raise NameError('foo')

self.assertFalse(entered)
with self.assertRaisesRegex(NameError, 'foo'):
await test()
self.assertFalse(entered)

@_async_test
async def test_decorating_method(self):

@asynccontextmanager
async def context():
yield


class Test(object):

@context()
async def method(self, a, b, c=None):
self.a = a
self.b = b
self.c = c

# these tests are for argument passing when used as a decorator
test = Test()
await test.method(1, 2)
self.assertEqual(test.a, 1)
self.assertEqual(test.b, 2)
self.assertEqual(test.c, None)

test = Test()
await test.method('a', 'b', 'c')
self.assertEqual(test.a, 'a')
self.assertEqual(test.b, 'b')
self.assertEqual(test.c, 'c')

test = Test()
await test.method(a=1, b=2)
self.assertEqual(test.a, 1)
self.assertEqual(test.b, 2)


class AclosingTestCase(unittest.TestCase):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Added missing behavior to :func:`contextlib.asynccontextmanager` to match
:func:`contextlib.contextmanager` so decorated functions can themselves be
decorators.

0 comments on commit 86b833b

Please sign in to comment.