concurrency.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. from __future__ import annotations
  2. import functools
  3. import warnings
  4. from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator
  5. from typing import ParamSpec, TypeVar
  6. import anyio.to_thread
  7. P = ParamSpec("P")
  8. T = TypeVar("T")
  9. async def run_until_first_complete(*args: tuple[Callable, dict]) -> None: # type: ignore[type-arg]
  10. warnings.warn(
  11. "run_until_first_complete is deprecated and will be removed in a future version.",
  12. DeprecationWarning,
  13. )
  14. async with anyio.create_task_group() as task_group:
  15. async def run(func: Callable[[], Coroutine]) -> None: # type: ignore[type-arg]
  16. await func()
  17. task_group.cancel_scope.cancel()
  18. for func, kwargs in args:
  19. task_group.start_soon(run, functools.partial(func, **kwargs))
  20. async def run_in_threadpool(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
  21. func = functools.partial(func, *args, **kwargs)
  22. return await anyio.to_thread.run_sync(func)
  23. class _StopIteration(Exception):
  24. pass
  25. def _next(iterator: Iterator[T]) -> T:
  26. # We can't raise `StopIteration` from within the threadpool iterator
  27. # and catch it outside that context, so we coerce them into a different
  28. # exception type.
  29. try:
  30. return next(iterator)
  31. except StopIteration:
  32. raise _StopIteration
  33. async def iterate_in_threadpool(
  34. iterator: Iterable[T],
  35. ) -> AsyncIterator[T]:
  36. as_iterator = iter(iterator)
  37. while True:
  38. try:
  39. yield await anyio.to_thread.run_sync(_next, as_iterator)
  40. except _StopIteration:
  41. break