44from collections import deque
55from functools import wraps
66
7- __all__ = ["contextmanager " , "closing " , "AbstractContextManager " ,
8- "ContextDecorator " , "ExitStack " , "redirect_stdout " ,
9- "redirect_stderr" , "suppress" ]
7+ __all__ = ["asynccontextmanager " , "contextmanager " , "closing " ,
8+ "AbstractContextManager " , "ContextDecorator " , "ExitStack " ,
9+ "redirect_stdout" , " redirect_stderr""suppress" ]
1010
1111
1212class AbstractContextManager (abc .ABC ):
@@ -54,8 +54,8 @@ def inner(*args, **kwds):
5454 return inner
5555
5656
57- class _GeneratorContextManager ( ContextDecorator , AbstractContextManager ) :
58- """Helper for @contextmanager decorator ."""
57+ class _GeneratorContextManagerBase :
58+ """Shared functionality for @contextmanager and @asynccontextmanager ."""
5959
6060 def __init__ (self , func , args , kwds ):
6161 self .gen = func (* args , ** kwds )
@@ -71,6 +71,12 @@ def __init__(self, func, args, kwds):
7171 # for the class instead.
7272 # See http://bugs.python.org/issue19404 for more details.
7373
74+
75+ class _GeneratorContextManager (_GeneratorContextManagerBase ,
76+ AbstractContextManager ,
77+ ContextDecorator ):
78+ """Helper for @contextmanager decorator."""
79+
7480 def _recreate_cm (self ):
7581 # _GCM instances are one-shot context managers, so the
7682 # CM must be recreated each time a decorated function is
@@ -121,12 +127,61 @@ def __exit__(self, type, value, traceback):
121127 # fixes the impedance mismatch between the throw() protocol
122128 # and the __exit__() protocol.
123129 #
130+ # This cannot use 'except BaseException as exc' (as in the
131+ # async implementation) to maintain compatibility with
132+ # Python 2, where old-style class exceptions are not caught
133+ # by 'except BaseException'.
124134 if sys .exc_info ()[1 ] is value :
125135 return False
126136 raise
127137 raise RuntimeError ("generator didn't stop after throw()" )
128138
129139
140+ class _AsyncGeneratorContextManager (_GeneratorContextManagerBase ):
141+ """Helper for @asynccontextmanager."""
142+
143+ async def __aenter__ (self ):
144+ try :
145+ return await self .gen .__anext__ ()
146+ except StopAsyncIteration :
147+ raise RuntimeError ("generator didn't yield" ) from None
148+
149+ async def __aexit__ (self , typ , value , traceback ):
150+ if typ is None :
151+ try :
152+ await self .gen .__anext__ ()
153+ except StopAsyncIteration :
154+ return
155+ else :
156+ raise RuntimeError ("generator didn't stop" )
157+ else :
158+ if value is None :
159+ value = typ ()
160+ # See _GeneratorContextManager.__exit__ for comments on subtleties
161+ # in this implementation
162+ try :
163+ await self .gen .athrow (typ , value , traceback )
164+ raise RuntimeError ("generator didn't stop after throw()" )
165+ except StopAsyncIteration as exc :
166+ return exc is not value
167+ except RuntimeError as exc :
168+ if exc is value :
169+ return False
170+ # Avoid suppressing if a StopIteration exception
171+ # was passed to throw() and later wrapped into a RuntimeError
172+ # (see PEP 479 for sync generators; async generators also
173+ # have this behavior). But do this only if the exception wrapped
174+ # by the RuntimeError is actully Stop(Async)Iteration (see
175+ # issue29692).
176+ if isinstance (value , (StopIteration , StopAsyncIteration )):
177+ if exc .__cause__ is value :
178+ return False
179+ raise
180+ except BaseException as exc :
181+ if exc is not value :
182+ raise
183+
184+
130185def contextmanager (func ):
131186 """@contextmanager decorator.
132187
@@ -153,14 +208,46 @@ def some_generator(<arguments>):
153208 <body>
154209 finally:
155210 <cleanup>
156-
157211 """
158212 @wraps (func )
159213 def helper (* args , ** kwds ):
160214 return _GeneratorContextManager (func , args , kwds )
161215 return helper
162216
163217
218+ def asynccontextmanager (func ):
219+ """@asynccontextmanager decorator.
220+
221+ Typical usage:
222+
223+ @asynccontextmanager
224+ async def some_async_generator(<arguments>):
225+ <setup>
226+ try:
227+ yield <value>
228+ finally:
229+ <cleanup>
230+
231+ This makes this:
232+
233+ async with some_async_generator(<arguments>) as <variable>:
234+ <body>
235+
236+ equivalent to this:
237+
238+ <setup>
239+ try:
240+ <variable> = <value>
241+ <body>
242+ finally:
243+ <cleanup>
244+ """
245+ @wraps (func )
246+ def helper (* args , ** kwds ):
247+ return _AsyncGeneratorContextManager (func , args , kwds )
248+ return helper
249+
250+
164251class closing (AbstractContextManager ):
165252 """Context to automatically close something at the end of a block.
166253
0 commit comments