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