You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
1.8 KiB
62 lines
1.8 KiB
import functools
|
|
import typing
|
|
from typing import Any, AsyncGenerator, Iterator
|
|
|
|
import anyio
|
|
|
|
try:
|
|
import contextvars # Python 3.7+ only or via contextvars backport.
|
|
except ImportError: # pragma: no cover
|
|
contextvars = None # type: ignore
|
|
|
|
|
|
T = typing.TypeVar("T")
|
|
|
|
|
|
async def run_until_first_complete(*args: typing.Tuple[typing.Callable, dict]) -> None:
|
|
async with anyio.create_task_group() as task_group:
|
|
|
|
async def run(func: typing.Callable[[], typing.Coroutine]) -> None:
|
|
await func()
|
|
task_group.cancel_scope.cancel()
|
|
|
|
for func, kwargs in args:
|
|
task_group.start_soon(run, functools.partial(func, **kwargs))
|
|
|
|
|
|
async def run_in_threadpool(
|
|
func: typing.Callable[..., T], *args: typing.Any, **kwargs: typing.Any
|
|
) -> T:
|
|
if contextvars is not None: # pragma: no cover
|
|
# Ensure we run in the same context
|
|
child = functools.partial(func, *args, **kwargs)
|
|
context = contextvars.copy_context()
|
|
func = context.run
|
|
args = (child,)
|
|
elif kwargs: # pragma: no cover
|
|
# run_sync doesn't accept 'kwargs', so bind them in here
|
|
func = functools.partial(func, **kwargs)
|
|
return await anyio.to_thread.run_sync(func, *args)
|
|
|
|
|
|
class _StopIteration(Exception):
|
|
pass
|
|
|
|
|
|
def _next(iterator: Iterator) -> Any:
|
|
# We can't raise `StopIteration` from within the threadpool iterator
|
|
# and catch it outside that context, so we coerce them into a different
|
|
# exception type.
|
|
try:
|
|
return next(iterator)
|
|
except StopIteration:
|
|
raise _StopIteration
|
|
|
|
|
|
async def iterate_in_threadpool(iterator: Iterator) -> AsyncGenerator:
|
|
while True:
|
|
try:
|
|
yield await anyio.to_thread.run_sync(_next, iterator)
|
|
except _StopIteration:
|
|
break
|