| from contextlib import ContextDecorator |
| from functools import wraps |
| |
| try: |
| from contextlib import asynccontextmanager |
| except ImportError: |
| asynccontextmanager = None |
| |
| |
| class _AsyncGeneratorContextManager(ContextDecorator): |
| def __init__(self, func, args, kwds): |
| self.gen = func(*args, **kwds) |
| self.func, self.args, self.kwds = func, args, kwds |
| self.__doc__ = func.__doc__ |
| |
| async def __aenter__(self): |
| return await self.gen.__anext__() |
| |
| async def __aexit__(self, typ, value, traceback): |
| if typ is not None: |
| await self.gen.athrow(typ, value, traceback) |
| |
| |
| def _asynccontextmanager(func): |
| @wraps(func) |
| def helper(*args, **kwds): |
| return _AsyncGeneratorContextManager(func, args, kwds) |
| |
| return helper |
| |
| |
| if asynccontextmanager is None: |
| asynccontextmanager = _asynccontextmanager |
