1from __future__ import annotations
2
3from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union, overload
4
5from typing_extensions import TypeAlias, TypeGuard
6
7if TYPE_CHECKING:
8 from stdlibx.matchtools._types import Operation
9
10T = TypeVar("T")
11U = TypeVar("U")
12
13T1 = TypeVar("T1")
14T2 = TypeVar("T2")
15T3 = TypeVar("T3")
16T4 = TypeVar("T4")
17T5 = TypeVar("T5")
18T6 = TypeVar("T6")
19T7 = TypeVar("T7")
20T8 = TypeVar("T8")
21T9 = TypeVar("T9")
22
23U1 = TypeVar("U1")
24U2 = TypeVar("U2")
25U3 = TypeVar("U3")
26U4 = TypeVar("U4")
27U5 = TypeVar("U5")
28U6 = TypeVar("U6")
29U7 = TypeVar("U7")
30U8 = TypeVar("U8")
31U9 = TypeVar("U9")
32
33
34Case: TypeAlias = tuple[Callable[[Any], TypeGuard[T]], Callable[[T], U]]
35
36
[docs]
37def type_is(type_: type[T], func: Callable[[T], U]) -> Case[T, U]:
38 return (lambda val: isinstance(val, type_), func) # type: ignore
39
40
41@overload
42def cases(
43 a: Case[T1, U1], b: Case[T2, U2], /
44) -> Operation[Union[T1, T2], Union[U1, U2]]: ...
45
46
47@overload
48def cases(
49 a: Case[T1, U1], b: Case[T2, U2], c: Case[T3, U3], /
50) -> Operation[Union[T1, T2, T3], Union[U1, U2, U3]]: ...
51
52
53@overload
54def cases(
55 a: Case[T1, U1], b: Case[T2, U2], c: Case[T3, U3], d: Case[T4, U4], /
56) -> Operation[Union[T1, T2, T3, T4], Union[U1, U2, U3, U4]]: ...
57
58
59@overload
60def cases(
61 a: Case[T1, U1],
62 b: Case[T2, U2],
63 c: Case[T3, U3],
64 d: Case[T4, U4],
65 e: Case[T5, U5],
66 /,
67) -> Operation[Union[T1, T2, T3, T4, T5], Union[U1, U2, U3, U4, U5]]: ...
68
69
70@overload
71def cases(
72 a: Case[T1, U1],
73 b: Case[T2, U2],
74 c: Case[T3, U3],
75 d: Case[T4, U4],
76 e: Case[T5, U5],
77 f: Case[T6, U6],
78 /,
79) -> Operation[Union[T1, T2, T3, T4, T5, T6], Union[U1, U2, U3, U4, U5, U6]]: ...
80
81
82@overload
83def cases(
84 a: Case[T1, U1],
85 b: Case[T2, U2],
86 c: Case[T3, U3],
87 d: Case[T4, U4],
88 e: Case[T5, U5],
89 f: Case[T6, U6],
90 g: Case[T7, U7],
91 /,
92) -> Operation[
93 Union[T1, T2, T3, T4, T5, T6, T7],
94 Union[U1, U2, U3, U4, U5, U6, U7],
95]: ...
96
97
98@overload
99def cases(
100 a: Case[T1, U1],
101 b: Case[T2, U2],
102 c: Case[T3, U3],
103 d: Case[T4, U4],
104 e: Case[T5, U5],
105 f: Case[T6, U6],
106 g: Case[T7, U7],
107 h: Case[T8, U8],
108 /,
109) -> Operation[
110 Union[T1, T2, T3, T4, T5, T6, T7, T8],
111 Union[U1, U2, U3, U4, U5, U6, U7, U8],
112]: ...
113
114
115@overload
116def cases(
117 a: Case[T1, U1],
118 b: Case[T2, U2],
119 c: Case[T3, U3],
120 d: Case[T4, U4],
121 e: Case[T5, U5],
122 f: Case[T6, U6],
123 g: Case[T7, U7],
124 h: Case[T8, U8],
125 i: Case[T9, U9],
126 /,
127) -> Operation[
128 Union[T1, T2, T3, T4, T5, T6, T7, T8, T9],
129 Union[U1, U2, U3, U4, U5, U6, U7, U8, U9],
130]: ...
131
132
[docs]
133def cases(*conditions: Case[Any, Any]) -> Operation[Any, Any]:
134 def _execute(value: Any) -> Any:
135 for condition, op in conditions:
136 if condition(value) is True:
137 return op(value)
138
139 msg = "Failed to handle all cases, make sure to include a catch-all case"
140 raise RuntimeError(msg)
141
142 return _execute