1from __future__ import annotations
2
3import functools
4from typing import TYPE_CHECKING, Any, Iterable, TypeVar, Union, overload
5
6from stdlibx.result._result import error, is_err, ok
7
8if TYPE_CHECKING:
9 from stdlibx.result.types import Result
10
11T = TypeVar("T")
12E = TypeVar("E")
13
14T1 = TypeVar("T1")
15T2 = TypeVar("T2")
16T3 = TypeVar("T3")
17T4 = TypeVar("T4")
18T5 = TypeVar("T5")
19T6 = TypeVar("T6")
20T7 = TypeVar("T7")
21T8 = TypeVar("T8")
22T9 = TypeVar("T9")
23
24E1 = TypeVar("E1")
25E2 = TypeVar("E2")
26E3 = TypeVar("E3")
27E4 = TypeVar("E4")
28E5 = TypeVar("E5")
29E6 = TypeVar("E6")
30E7 = TypeVar("E7")
31E8 = TypeVar("E8")
32E9 = TypeVar("E9")
33
34
35@overload
36def collect(
37 a: Result[T1, E1],
38 b: Result[T2, E2],
39 /,
40) -> Result[tuple[T1, T2], Union[E1, E2]]: ...
41
42
43@overload
44def collect(
45 a: Result[T1, E1],
46 b: Result[T2, E2],
47 c: Result[T3, E3],
48 /,
49) -> Result[tuple[T1, T2, T3], Union[E1, E2, E3]]: ...
50
51
52@overload
53def collect(
54 a: Result[T1, E1],
55 b: Result[T2, E2],
56 c: Result[T3, E3],
57 d: Result[T4, E4],
58 /,
59) -> Result[tuple[T1, T2, T3, T4], Union[E1, E2, E3, E4]]: ...
60
61
62@overload
63def collect(
64 a: Result[T1, E1],
65 b: Result[T2, E2],
66 c: Result[T3, E3],
67 d: Result[T4, E4],
68 e: Result[T5, E5],
69 /,
70) -> Result[tuple[T1, T2, T3, T4, T5], Union[E1, E2, E3, E4, E5]]: ...
71
72
73@overload
74def collect(
75 a: Result[T1, E1],
76 b: Result[T2, E2],
77 c: Result[T3, E3],
78 d: Result[T4, E4],
79 e: Result[T5, E5],
80 f: Result[T6, E6],
81 /,
82) -> Result[tuple[T1, T2, T3, T4, T5, T6], Union[E1, E2, E3, E4, E5, E6]]: ...
83
84
85@overload
86def collect(
87 a: Result[T1, E1],
88 b: Result[T2, E2],
89 c: Result[T3, E3],
90 d: Result[T4, E4],
91 e: Result[T5, E5],
92 f: Result[T6, E6],
93 g: Result[T7, E7],
94 /,
95) -> Result[tuple[T1, T2, T3, T4, T5, T6, T7], Union[E1, E2, E3, E4, E5, E6, E7]]: ...
96
97
98@overload
99def collect(
100 a: Result[T1, E1],
101 b: Result[T2, E2],
102 c: Result[T3, E3],
103 d: Result[T4, E4],
104 e: Result[T5, E5],
105 f: Result[T6, E6],
106 g: Result[T7, E7],
107 h: Result[T8, E8],
108 /,
109) -> Result[
110 tuple[T1, T2, T3, T4, T5, T6, T7, T8],
111 Union[E1, E2, E3, E4, E5, E6, E7, E8],
112]: ...
113
114
115@overload
116def collect(
117 a: Result[T1, E1],
118 b: Result[T2, E2],
119 c: Result[T3, E3],
120 d: Result[T4, E4],
121 e: Result[T5, E5],
122 f: Result[T6, E6],
123 g: Result[T7, E7],
124 h: Result[T8, E8],
125 i: Result[T9, E9],
126 /,
127) -> Result[
128 tuple[T1, T2, T3, T4, T5, T6, T7, T8, T9],
129 Union[E1, E2, E3, E4, E5, E6, E7, E8, E9],
130]: ...
131
132
[docs]
133def collect(
134 initial: Result[Any, Any],
135 *others: Result[Any, Any],
136) -> Result[tuple[Any, ...], Any]:
137 def _combine(
138 a: Result[tuple[Any, ...], Any],
139 b: Result[Any, Any],
140 ) -> Result[tuple[Any, ...], Any]:
141 if is_err(a):
142 return error(a.error)
143 elif is_err(b):
144 return error(b.error)
145 else:
146 return ok(((*a.value, b.value))) # type: ignore
147
148 return functools.reduce(_combine, [initial, *others], ok(()))
149
150
[docs]
151def collect_all(iterable: Iterable[Result[T, E]]) -> Result[tuple[T, ...], E]:
152 def _combine(
153 a: Result[tuple[T, ...], E],
154 b: Result[T, E],
155 ) -> Result[tuple[T, ...], E]:
156 if is_err(a):
157 return error(a.error)
158 elif is_err(b):
159 return error(b.error)
160 else:
161 return ok(((*a.value, b.value))) # type: ignore
162
163 return functools.reduce(_combine, iterable, ok(()))