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