8000 bpo-40816 Add AsyncContextDecorator class (GH-20516) · python/cpython@178695b · GitHub
[go: up one dir, main page]

Skip to content

Commit 178695b

Browse files
heckad1st1
andauthored
bpo-40816 Add AsyncContextDecorator class (GH-20516)
Co-authored-by: Yury Selivanov <yury@edgedb.com>
1 parent 048a356 commit 178695b

File tree

4 files changed

+114
-1
lines changed

4 files changed

+114
-1
lines changed

Doc/library/contextlib.rst

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,31 @@ Functions and classes provided:
126126

127127
.. versionadded:: 3.7
128128

129+
Context managers defined with :func:`asynccontextmanager` can be used
130+
either as decorators or with :keyword:`async with` statements::
131+
132+
import time
133+
134+
async def timeit():
135+
now = time.monotonic()
136+
try:
137+
yield
138+
finally:
139+
print(f'it took {time.monotonic() - now}s to run')
140+
141+
@timeit()
142+
async def main():
143+
# ... async code ...
144+
145+
When used as a decorator, a new generator instance is implicitly created on
146+
each function call. This allows the otherwise "one-shot" context managers
147+
created by :func:`asynccontextmanager` to meet the requirement that context
148+
managers support multiple invocations in order to be used as decorators.
149+
150+
.. versionchanged:: 3.10
151+
Async context managers created with :func:`asynccontextmanager` can
152+
be used as decorators.
153+
129154

130155
.. function:: closing(thing)
131156

@@ -384,6 +409,43 @@ Functions and classes provided:
384409
.. versionadded:: 3.2
385410

386411

412+
.. class:: AsyncContextManager
413+
414+
Similar as ContextManger only for async
415+
416+
Example of ``ContextDecorator``::
417+
418+
from asyncio import run
419+
from contextlib import AsyncContextDecorator
420+
421+
class mycontext(AsyncContextDecorator):
422+
async def __aenter__(self):
423+
print('Starting')
424+
return self
425+
426+
async def __aexit__(self, *exc):
427+
print('Finishing')
428+
return False
429+
430+
>>> @mycontext()
431+
... async def function():
432+
... print('The bit in the middle')
433+
...
434+
>>> run(function())
435+
Starting
436+
The bit in the middle
437+
Finishing
438+
439+
>>> async def function():
440+
... async with mycontext():
441+
... print('The bit in the middle')
442+
...
443+
>>> run(function())
444+
Starting
445+
The bit in the middle
446+
Finishing
447+
448+
387449
.. class:: ExitStack()
388450

389451
A context manager that is designed to make it easy to programmatically

Lib/contextlib.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,22 @@ def inner(*args, **kwds):
8080
return inner
8181

8282

83+
class AsyncContextDecorator(object):
84+
"A base class or mixin that enables async context managers to work as decorators."
85+
86+
def _recreate_cm(self):
87+
"""Return a recreated instance of self.
88+
"""
89+
return self
90+
91+
def __call__(self, func):
92+
@wraps(func)
93+
async def inner(*args, **kwds):
94+
async with self._recreate_cm():
95+
return await func(*args, **kwds)
96+
return inner
97+
98+
8399
class _GeneratorContextManagerBase:
84100
"""Shared functionality for @contextmanager and @asynccontextmanager."""
85101

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

168184

169185
class _AsyncGeneratorContextManager(_GeneratorContextManagerBase,
170-
AbstractAsyncContextManager):
186+
AbstractAsyncContextManager,
187+
AsyncContextDecorator):
171188
"""Helper for @asynccontextmanager."""
172189

190+
def _recreate_cm(self):
191+
# _AGCM instances are one-shot context managers, so the
192+
# ACM must be recreated each time a decorated function is
193+
# called
194+
return self.__class__(self.func, self.args, self.kwds)
195+
173196
async def __aenter__(self):
174197
try:
175198
return await self.gen.__anext__()

Lib/test/test_contextlib_async.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,33 @@ async def woohoo(self, func, args, kwds):
278278
async with woohoo(self=11, func=22, args=33, kwds=44) as target:
279279
self.assertEqual(target, (11, 22, 33, 44))
280280

281+
@_async_test
282+
async def test_recursive(self):
283+
depth = 0
284+
ncols = 0
285+
286+
@asynccontextmanager
287+
async def woohoo():
288+
nonlocal ncols
289+
ncols += 1
290+
291+
nonlocal depth
292+
before = depth
293+
depth += 1
294+
yield
295+
depth -= 1
296+
self.assertEqual(depth, before)
297+
298+
@woohoo()
299+
async def recursive():
300+
if depth < 10:
301+
await recursive()
302+
303+
await recursive()
304+
305+
self.assertEqual(ncols, 10)
306+
self.assertEqual(depth, 0)
307+
281308

282309
class AclosingTestCase(unittest.TestCase):
283310

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add AsyncContextDecorator to contextlib to support async context manager as a decorator.

0 commit comments

Comments
 (0)
0