-
-
Notifications
You must be signed in to change notification settings - Fork 6.1k
/
concurrency.py
25 lines (21 loc) · 880 Bytes
/
concurrency.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from contextlib import AsyncExitStack as AsyncExitStack # noqa
from contextlib import asynccontextmanager as asynccontextmanager
from typing import AsyncGenerator, ContextManager, TypeVar
from starlette.concurrency import iterate_in_threadpool as iterate_in_threadpool # noqa
from starlette.concurrency import run_in_threadpool as run_in_threadpool # noqa
from starlette.concurrency import ( # noqa
run_until_first_complete as run_until_first_complete,
)
_T = TypeVar("_T")
@asynccontextmanager
async def contextmanager_in_threadpool(
cm: ContextManager[_T],
) -> AsyncGenerator[_T, None]:
try:
yield await run_in_threadpool(cm.__enter__)
except Exception as e:
ok: bool = await run_in_threadpool(cm.__exit__, type(e), e, None)
if not ok:
raise e
else:
await run_in_threadpool(cm.__exit__, None, None, None)