1from __future__ import annotations
2
3import functools
4from typing import TYPE_CHECKING, Any, Iterable, TypeVar, overload
5
6from stdlibx.option._option import is_some, nothing, some
7
8if TYPE_CHECKING:
9 from stdlibx.option.types import Option
10
11T = TypeVar("T")
12T1 = TypeVar("T1")
13T2 = TypeVar("T2")
14T3 = TypeVar("T3")
15T4 = TypeVar("T4")
16T5 = TypeVar("T5")
17T6 = TypeVar("T6")
18T7 = TypeVar("T7")
19T8 = TypeVar("T8")
20T9 = TypeVar("T9")
21
22
23@overload
24def collect(a: Option[T1], b: Option[T2], /) -> Option[tuple[T1, T2]]: ...
25
26
27@overload
28def collect(
29 a: Option[T1], b: Option[T2], c: Option[T3], /
30) -> Option[tuple[T1, T2, T3]]: ...
31
32
33@overload
34def collect(
35 a: Option[T1], b: Option[T2], c: Option[T3], d: Option[T4], /
36) -> Option[tuple[T1, T2, T3, T4]]: ...
37
38
39@overload
40def collect(
41 a: Option[T1], b: Option[T2], c: Option[T3], d: Option[T4], e: Option[T5], /
42) -> Option[tuple[T1, T2, T3, T4, T5]]: ...
43
44
45@overload
46def collect(
47 a: Option[T1],
48 b: Option[T2],
49 c: Option[T3],
50 d: Option[T4],
51 e: Option[T5],
52 f: Option[T6],
53 /,
54) -> Option[tuple[T1, T2, T3, T4, T5, T6]]: ...
55
56
57@overload
58def collect(
59 a: Option[T1],
60 b: Option[T2],
61 c: Option[T3],
62 d: Option[T4],
63 e: Option[T5],
64 f: Option[T6],
65 g: Option[T7],
66 /,
67) -> Option[tuple[T1, T2, T3, T4, T5, T6, T7]]: ...
68
69
70@overload
71def collect(
72 a: Option[T1],
73 b: Option[T2],
74 c: Option[T3],
75 d: Option[T4],
76 e: Option[T5],
77 f: Option[T6],
78 g: Option[T7],
79 h: Option[T8],
80 /,
81) -> Option[tuple[T1, T2, T3, T4, T5, T6, T7, T8]]: ...
82
83
84@overload
85def collect(
86 a: Option[T1],
87 b: Option[T2],
88 c: Option[T3],
89 d: Option[T4],
90 e: Option[T5],
91 f: Option[T6],
92 g: Option[T7],
93 h: Option[T8],
94 i: Option[T9],
95 /,
96) -> Option[tuple[T1, T2, T3, T4, T5, T6, T7, T8, T9]]: ...
97
98
[docs]
99def collect(initial: Option[Any], *others: Option[Any]) -> Option[tuple[Any, ...]]:
100 def _combine(a: Option[tuple[Any, ...]], b: Option[Any]) -> Option[tuple[Any, ...]]:
101 if is_some(a) and is_some(b):
102 return some(((*a.value, b.value)))
103 return nothing()
104
105 return functools.reduce(_combine, [initial, *others], some(()))
106
107
[docs]
108def collect_all(iterable: Iterable[Option[T]]) -> Option[tuple[T, ...]]:
109 def _combine(a: Option[tuple[T, ...]], b: Option[T]) -> Option[tuple[T, ...]]:
110 if is_some(a) and is_some(b):
111 return some(((*a.value, b.value)))
112 return nothing()
113
114 return functools.reduce(_combine, iterable, some(()))