Add extra overloads for starpipe

dbrattli · Jan 24, 2024 · a24eb37 · a24eb37
1 parent 3cfd9e8
commit a24eb37
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/expression/core/pipe.py b/expression/core/pipe.py
@@ -33,6 +33,7 @@
 _X = TypeVarTuple("_X")
 _Y = TypeVarTuple("_Y")
 _Z = TypeVarTuple("_Z")
+_K = TypeVarTuple("_K")
 
 
 @overload
@@ -223,6 +224,31 @@ def starpipe(
     ...
 
 
+@overload
+def starpipe(
+    __args: tuple[*_P],
+    __fn1: Callable[[*_P], tuple[*_Q]],
+    __fn2: Callable[[*_Q], tuple[*_X]],
+    __fn3: Callable[[*_X], tuple[*_Y]],
+    __fn4: Callable[[*_Y], tuple[*_Z]],
+    __fn5: Callable[[*_Z], _B],
+) -> _B:
+    ...
+
+
+@overload
+def starpipe(
+    __args: tuple[*_P],
+    __fn1: Callable[[*_P], tuple[*_Q]],
+    __fn2: Callable[[*_Q], tuple[*_X]],
+    __fn3: Callable[[*_X], tuple[*_Y]],
+    __fn4: Callable[[*_Y], tuple[*_Z]],
+    __fn5: Callable[[*_Z], tuple[*_K]],
+    __fn6: Callable[[*_K], _B],
+) -> _B:
+    ...
+
+
 def starpipe(__args: Any, *__fns: Callable[..., Any]) -> Any:
     """Functional pipe_n (`||>`, `||>`, `|||>`, etc).