forked from fastapi/fastapi
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathconcurrency.py
More file actions
51 lines (41 loc) · 1.68 KB
/
Copy pathconcurrency.py
File metadata and controls
51 lines (41 loc) · 1.68 KB
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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
from typing import Any, Callable
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,
)
asynccontextmanager_error_message = """
FastAPI's contextmanager_in_threadpool require Python 3.7 or above,
or the backport for Python 3.6, installed with:
pip install async-generator
"""
def _fake_asynccontextmanager(func: Callable[..., Any]) -> Callable[..., Any]:
def raiser(*args: Any, **kwargs: Any) -> Any:
raise RuntimeError(asynccontextmanager_error_message)
return raiser
try:
from contextlib import asynccontextmanager as asynccontextmanager # type: ignore
except ImportError:
try:
from async_generator import ( # type: ignore # isort: skip
asynccontextmanager as asynccontextmanager,
)
except ImportError: # pragma: no cover
asynccontextmanager = _fake_asynccontextmanager
try:
from contextlib import AsyncExitStack as AsyncExitStack # type: ignore
except ImportError:
try:
from async_exit_stack import AsyncExitStack as AsyncExitStack # type: ignore
except ImportError: # pragma: no cover
AsyncExitStack = None # type: ignore
@asynccontextmanager # type: ignore
async def contextmanager_in_threadpool(cm: Any) -> Any:
try:
yield await run_in_threadpool(cm.__enter__)
except Exception as e:
ok = 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)