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