Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/_core/_concurrency_utils.py: 53%

1from __future__ import annotations 

2 

3from inspect import iscoroutine 

4 

5__all__ = ("amap", "as_completed", "gather") 

6 

7from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Iterable 

8from contextlib import asynccontextmanager 

9from typing import TYPE_CHECKING, Any, TypeVar, overload 

10 

11from ..abc._tasks import get_coro_name 

12from ._exceptions import BrokenResourceError 

13from ._streams import create_memory_object_stream 

14from ._tasks import TaskHandle, create_task_group 

15 

16if TYPE_CHECKING: 

17 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream 

18 

19R = TypeVar("R") 

20S = TypeVar("S") 

21T = TypeVar("T") 

22U = TypeVar("U") 

23V = TypeVar("V") 

24W = TypeVar("W") 

25 

26 

27@overload 

28async def gather( 

29 coro1: Coroutine[Any, Any, R], coro2: Coroutine[Any, Any, S], / 

30) -> tuple[R, S]: ... 

31 

32 

33@overload 

34async def gather( 

35 coro1: Coroutine[Any, Any, R], 

36 coro2: Coroutine[Any, Any, S], 

37 coro3: Coroutine[Any, Any, T], 

38 /, 

39) -> tuple[R, S, T]: ... 

40 

41 

42@overload 

43async def gather( 

44 coro1: Coroutine[Any, Any, R], 

45 coro2: Coroutine[Any, Any, S], 

46 coro3: Coroutine[Any, Any, T], 

47 coro4: Coroutine[Any, Any, U], 

48 /, 

49) -> tuple[R, S, T, U]: ... 

50 

51 

52@overload 

53async def gather( 

54 coro1: Coroutine[Any, Any, R], 

55 coro2: Coroutine[Any, Any, S], 

56 coro3: Coroutine[Any, Any, T], 

57 coro4: Coroutine[Any, Any, U], 

58 coro5: Coroutine[Any, Any, V], 

59 /, 

60) -> tuple[R, S, T, U, V]: ... 

61 

62 

63@overload 

64async def gather( 

65 coro1: Coroutine[Any, Any, R], 

66 coro2: Coroutine[Any, Any, S], 

67 coro3: Coroutine[Any, Any, T], 

68 coro4: Coroutine[Any, Any, U], 

69 coro5: Coroutine[Any, Any, V], 

70 coro6: Coroutine[Any, Any, W], 

71 /, 

72) -> tuple[R, S, T, U, V, W]: ... 

73 

74 

75# handle arbitrary length if awaitables are all of the same type 

76@overload 

77async def gather(*coros: Coroutine[Any, Any, R]) -> tuple[R, ...]: ... 

78 

79 

80async def gather(*coros: Coroutine[Any, Any, Any]) -> tuple[Any, ...]: 

81 """ 

82 Run coroutines concurrently in a task group. The order of result values corresponds 

83 to the order of coroutines passed. 

84 

85 :param coros: coroutine objects to run as tasks 

86 :return: task results for each argument in the same order as they were passed 

87 

88 """ 

89 handles: list[TaskHandle[Any]] = [] 

90 async with create_task_group() as tg: 

91 handles.extend(tg.create_task(coro) for coro in coros) 

92 

93 return tuple(h.return_value for h in handles) 

94 

95 

96@asynccontextmanager 

97async def as_completed( 

98 *awaitables: Awaitable[R], 

99) -> AsyncGenerator[MemoryObjectReceiveStream[TaskHandle[R]]]: 

100 """ 

101 Run awaitable objects concurrently in a task group, returning an iterator which can 

102 be used to get finished task handles in the order they complete. 

103 

104 :param awaitables: awaitable objects to run as tasks 

105 :return: MemoryObjectReceiveStream for iterating over task handles as tasks complete. 

106 

107 """ 

108 if not awaitables: 

109 raise ValueError("as_completed() takes at least one awaitable") 

110 

111 send, recv = create_memory_object_stream[TaskHandle[R]](len(awaitables)) 

112 task_handles: list[TaskHandle[R]] = [] 

113 

114 async def runner( 

115 awaitable: Awaitable[R], 

116 index: int, 

117 _send: MemoryObjectSendStream[TaskHandle[R]], 

118 ) -> R: 

119 async with _send: 

120 try: 

121 return await awaitable 

122 finally: 

123 try: 

124 _send.send_nowait(task_handles[index]) 

125 except BrokenResourceError: 

126 pass 

127 

128 async with recv, create_task_group() as tg: 

129 async with send: 

130 for i, awaitable in enumerate(awaitables): 

131 name = get_coro_name(awaitable) if iscoroutine(awaitable) else None 

132 task_handles.append( 

133 tg.start_soon(runner, awaitable, i, send.clone(), name=name) 

134 ) 

135 try: 

136 yield recv 

137 finally: 

138 tg.cancel_scope.cancel() 

139 

140 

141async def amap( 

142 func: Callable[[T], Coroutine[Any, Any, R]], args: Iterable[T] 

143) -> list[R]: 

144 """ 

145 Run the given coroutine function concurrently for multiple argument values. 

146 

147 :param func: a coroutine function that takes a single argument 

148 :param args: a sequence of argument values to pass to ``func`` 

149 :return: task results for each argument in the same order as they were passed 

150 

151 """ 

152 handles: list[TaskHandle[R]] = [] 

153 async with create_task_group() as tg: 

154 handles.extend(tg.start_soon(func, arg) for arg in args) 

155 

156 return [h.return_value for h in handles]