1from asyncutils._internal import helpers as H, patch as P
2from asyncutils._internal.submodules import compete_all as __all__
3import asyncio as I, asyncutils as A
4from asyncio.staggered import staggered_race
5from sys import audit
[docs]
6async def first_completed(*C, ret_exc=False, timeout=None):
7 if not C: raise TypeError('asyncutils.compete.first_completed: pass in at least one coroutine')
8 audit('asyncutils.compete.first_completed/start', L := len(C)); t = tuple(A.new_eager_tasks(*C))
9 try:
10 async with I.timeout(timeout):
11 for F in (await I.wait(t, return_when='FIRST_COMPLETED'))[0]: return A.wrap_exc(e) if ret_exc and (e := F.exception()) else F.result()
12 finally: audit('asyncutils.compete.first_completed/end', L); await A.safe_cancel_batch(t)
[docs]
13async def race_with_callback(*C, winner=None, loser=None, timeout=None):
14 if not C: raise TypeError('asyncutils.compete.race_with_callback: pass in at least one coroutine')
15 audit('asyncutils.compete.race_with_callback/start', L := len(C)); d, p = await I.wait(A.new_eager_tasks(*C), return_when='FIRST_COMPLETED', timeout=timeout)
16 try:
17 if not d: return None
18 w = d.pop().result()
19 if winner is not None and I.iscoroutine(r := winner(w)): await r
20 return w
21 finally: audit('asyncutils.compete.race_with_callback/end', L); await A.safe_cancel_batch(p, callback=loser)
[docs]
22async def multi_winner_race_with_callback(*C, timeout, winner=None, loser=None, _=__import__('operator').methodcaller('result')): # noqa: B008
23 if not C: raise TypeError('asyncutils.compete.multi_winner_race_with_callback: pass in at least one coroutine')
24 audit('asyncutils.compete.multi_winner_race_with_callback/start', L := len(C)); d, p = await I.wait(A.new_eager_tasks(*C), timeout=timeout); d = map(_, d)
25 try:
26 if winner is None: return list(d)
27 async def g(a, /, _=winner, f=(r := []).append):
28 if I.iscoroutine(x := _(a)): await x
29 f(a)
30 await I.gather(*map(g, d)); return r
31 except A.CRITICAL: raise A.Critical
32 finally: audit('asyncutils.compete.multi_winner_race_with_callback/end', L); await A.safe_cancel_batch(p, callback=loser)
[docs]
33def convert_to_coro_iter(cfs, *, loop=None, skip_invalid=None, corocheck=I.iscoroutine, futwrap=I.wrap_future, handle_aiter=None, handle_iter=None, _c=H.check_methods):
34 if handle_iter is None: from asyncutils import to_list as handle_iter
35 if handle_aiter is None: from asyncutils import to_list as handle_aiter
36 if skip_invalid is None: from asyncutils.context import CONVERT_TO_CORO_ITER_DEFAULT_SKIP_INVALID as skip_invalid
37 for i in A.aiter_to_gen(cfs, loop=loop):
38 if corocheck(i): yield i; continue
39 try: i = futwrap(i, loop=loop) # noqa: PLW2901
40 except A.CRITICAL: raise A.Critical
41 except (AssertionError, TypeError):
42 if not _c(i, '__await__'):
43 if _c(i, '__aiter__'): yield handle_aiter(i)
44 elif _c(i, '__iter__'): yield handle_iter(i)
45 elif not skip_invalid: raise TypeError(f'asyncutils.compete.convert_to_coro_iter: invalid item in {cfs!r}: {i!r}') from None
46 continue
47 yield A.wrap_in_coro(i)
[docs]
48def enhanced_staggered_race(cfs, delay=None, *, loop=None): return staggered_race(map(lambda c: lambda: c, convert_to_coro_iter(cfs, loop=loop)), delay, loop=loop)
[docs]
49def enhanced_gather(it, return_exceptions=False, *, loop=None, _=I.gather): return _(*convert_to_coro_iter(it, loop=loop), return_exceptions=return_exceptions)
50P.patch_function_signatures((first_completed, '*coros, ret_exc=False, timeout=None, loop=None'), (race_with_callback, '*coros, winner=None, loser=None, timeout=None'), (multi_winner_race_with_callback, '*coros, timeout, winner=None, loser=None'), (convert_to_coro_iter, 'cfs, *, loop=None, skip_invalid=None, corocheck={0}, futwrap={0}, handle_aiter=None, handle_iter=None'), (enhanced_gather, 'it, return_exceptions=False, *, loop=None'))
51del H, P