1from contextlib import ContextDecorator 2from functools import wraps 3 4try: 5 from contextlib import asynccontextmanager 6except ImportError: 7 asynccontextmanager = None 8 9 10class _AsyncGeneratorContextManager(ContextDecorator): 11 def __init__(self, func, args, kwds): 12 self.gen = func(*args, **kwds) 13 self.func, self.args, self.kwds = func, args, kwds 14 self.__doc__ = func.__doc__ 15 16 async def __aenter__(self): 17 return await self.gen.__anext__() 18 19 async def __aexit__(self, typ, value, traceback): 20 if typ is not None: 21 await self.gen.athrow(typ, value, traceback) 22 23 24def _asynccontextmanager(func): 25 @wraps(func) 26 def helper(*args, **kwds): 27 return _AsyncGeneratorContextManager(func, args, kwds) 28 29 return helper 30 31 32if asynccontextmanager is None: 33 asynccontextmanager = _asynccontextmanager 34