fix: starpipe bugfix and type hints

expression/collections/block.py

@@ -23,7 +23,7 @@
 import functools
 import itertools
 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, get_args, overload
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, TypeVarTuple, get_args, overload
 
 
 if TYPE_CHECKING:
@@ -54,6 +54,7 @@
 _T2 = TypeVar("_T2")
 _T3 = TypeVar("_T3")
 _T4 = TypeVar("_T4")
+_P = TypeVarTuple("_P")
 
 
 class Block(
@@ -235,25 +236,7 @@ def map(self, mapping: Callable[[_TSource], _TResult]) -> Block[_TResult]:
         """
         return Block((*builtins.map(mapping, self),))
 
-    @overload
-    def starmap(self: Block[tuple[_T1, _T2]], mapping: Callable[[_T1, _T2], _TResult]) -> Block[_TResult]:
-        ...
-
-    @overload
-    def starmap(
-        self: Block[tuple[_T1, _T2, _T3]],
-        mapping: Callable[[_T1, _T2, _T3], _TResult],
-    ) -> Block[_TResult]:
-        ...
-
-    @overload
-    def starmap(
-        self: Block[tuple[_T1, _T2, _T3, _T4]],
-        mapping: Callable[[_T1, _T2, _T3, _T4], _TResult],
-    ) -> Block[_TResult]:
-        ...
-
-    def starmap(self: Block[Any], mapping: Callable[..., Any]) -> Block[Any]:
+    def starmap(self: Block[tuple[*_P]], mapping: Callable[[*_P], _TResult]) -> Block[_TResult]:
         """Starmap source sequence.
 
         Unpack arguments grouped as tuple elements. Builds a new collection
@@ -767,24 +750,7 @@ def reduce(
     return source.tail().fold(reduction, source.head())
 
 
-@overload
-def starmap(mapper: Callable[[_T1, _T2], _TResult]) -> Callable[[Block[tuple[_T1, _T2]]], Block[_TResult]]:
-    ...
-
-
-@overload
-def starmap(mapper: Callable[[_T1, _T2, _T3], _TResult]) -> Callable[[Block[tuple[_T1, _T2, _T3]]], Block[_TResult]]:
-    ...
-
-
-@overload
-def starmap(
-    mapper: Callable[[_T1, _T2, _T3, _T4], _TResult],
-) -> Callable[[Block[tuple[_T1, _T2, _T3, _T4]]], Block[_TResult]]:
-    ...
-
-
-def starmap(mapper: Callable[..., Any]) -> Callable[[Block[Any]], Block[Any]]:
+def starmap(mapper: Callable[[*_P], _TResult]) -> Callable[[Block[tuple[*_P]]], Block[_TResult]]:
     """Starmap source sequence.
 
     Unpack arguments grouped as tuple elements. Builds a new collection
@@ -798,7 +764,7 @@ def starmap(mapper: Callable[..., Any]) -> Callable[[Block[Any]], Block[Any]]:
         Partially applied map function.
     """
 
-    def mapper_(args: tuple[Any, ...]) -> Any:
+    def mapper_(args: tuple[*_P]) -> _TResult:
         return mapper(*args)
 
     return map(mapper_)

expression/core/compose.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from functools import reduce
-from typing import Any, TypeVar, overload
+from typing import Any, TypeVar, TypeVarTuple, overload
 
 
 _A = TypeVar("_A")
@@ -14,6 +14,12 @@
 _T = TypeVar("_T")
 _J = TypeVar("_J")
 
+_P = TypeVarTuple("_P")
+_Q = TypeVarTuple("_Q")
+_X = TypeVarTuple("_X")
+_Y = TypeVarTuple("_Y")
+_Z = TypeVarTuple("_Z")
+
 
 @overload
 def compose() -> Callable[[_A], _A]:
@@ -138,4 +144,44 @@ def _compose(source: Any) -> Any:
     return _compose
 
 
+@overload
+def starcompose() -> Callable[[Any], Any]:
+    ...
+
+
+@overload
+def starcompose(__fn1: Callable[[*_P], _A]) -> Callable[[*_P], _A]:
+    ...
+
+
+@overload
+def starcompose(__fn1: Callable[[*_P], tuple[*_Y]], __fn2: Callable[[*_Y], _B]) -> Callable[[*_P], _B]:
+    ...
+
+
+@overload
+def starcompose(
+    __fn1: Callable[[*_P], tuple[*_Y]], __fn2: Callable[[*_Y], tuple[*_Z]], __fn3: Callable[[*_Z], _C]
+) -> Callable[[*_P], _C]:
+    ...
+
+
+@overload
+def starcompose(
+    __fn1: Callable[[*_P], tuple[*_Y]],
+    __fn2: Callable[[*_Y], tuple[*_Z]],
+    __fn3: Callable[[*_Z], tuple[*_X]],
+    __fn4: Callable[[*_X], _D],
+) -> Callable[[*_P], _D]:
+    ...
+
+
+def starcompose(*fns: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    def _compose(source: Any) -> Any:
+        """Return a pipeline of composed functions."""
+        return reduce(lambda fields, f: f(*fields), fns, source)
+
+    return _compose
+
+
 __all__ = ["compose"]

expression/core/misc.py

@@ -1,7 +1,7 @@
 from collections.abc import Callable
 from typing import Any, TypeVar
 
-from typing_extensions import TypeVarTuple, Unpack
+from typing_extensions import TypeVarTuple
 
 
 _A = TypeVar("_A")
@@ -19,7 +19,7 @@ def identity(value: _A) -> _A:
     return value
 
 
-def starid(*value: Unpack[_P]) -> tuple[Unpack[_P]]:
+def starid(*value: *_P) -> tuple[*_P]:
     return value
 
 

expression/core/option.py

@@ -9,7 +9,7 @@
 
 import builtins
 from collections.abc import Callable, Generator, Iterable
-from typing import TYPE_CHECKING, Any, Literal, TypeGuard, TypeVar, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Literal, TypeGuard, TypeVar, TypeVarTuple, get_args, get_origin
 
 from .curry import curry_flip
 from .error import EffectError
@@ -32,6 +32,7 @@
 
 _T1 = TypeVar("_T1")
 _T2 = TypeVar("_T2")
+_P = TypeVarTuple("_P")
 
 
 @tagged_union(frozen=True, order=True)
@@ -104,6 +105,19 @@ def map2(self, mapper: Callable[[_TSource, _T2], _TResult], other: Option[_T2])
             case _:
                 return Nothing
 
+    def starmap(self: Option[tuple[*_P]], mapper: Callable[[*_P], _TResult]) -> Option[_TResult]:
+        """Starmap option.
+
+        Applies the mapper to the values if the option is Some,
+        otherwise returns `Nothing`. The tuple is unpacked before
+        applying the mapper.
+        """
+        match self:
+            case Option(tag="some", some=some):
+                return Some(mapper(*some))
+            case _:
+                return Nothing
+
     def bind(self, mapper: Callable[[_TSource], Option[_TResult]]) -> Option[_TResult]:
         """Bind option.
 
@@ -415,6 +429,11 @@ def map2(opt1: Option[_T1], opt2: Option[_T2], mapper: Callable[[_T1, _T2], _TRe
     return opt1.map2(mapper, opt2)
 
 
+@curry_flip(1)
+def starmap(option: Option[tuple[*_P]], mapper: Callable[[*_P], _TResult]) -> Option[_TResult]:
+    return option.starmap(mapper)
+
+
 def or_else(
     option: Option[_TSource],
     if_none: Option[_TSource],

expression/core/pipe.py

@@ -12,9 +12,9 @@
     >>> assert pipe(v, fn, gn) == gn(fn(v))
 """
 from collections.abc import Callable
-from typing import Any, TypeVar, overload
+from typing import Any, TypeVar, TypeVarTuple, cast, overload
 
-from .compose import compose
+from .compose import compose, starcompose
 from .misc import starid
 
 
@@ -28,6 +28,11 @@
 _H = TypeVar("_H")
 _T = TypeVar("_T")
 _J = TypeVar("_J")
+_P = TypeVarTuple("_P")
+_Q = TypeVarTuple("_Q")
+_X = TypeVarTuple("_X")
+_Y = TypeVarTuple("_Y")
+_Z = TypeVarTuple("_Z")
 
 
 @overload
@@ -187,20 +192,54 @@ def pipe3(__values: Any, *fns: Any) -> Any:
     return pipe(fns[0](__values[0])(__values[1])(__values[2]), *fns[1:]) if fns else __values
 
 
-def starpipe(args: tuple[Any, ...], *fns: Callable[..., Any]):
+@overload
+def starpipe(__args: tuple[*_P], __fn1: Callable[[*_P], _B]) -> _B:
+    ...
+
+
+@overload
+def starpipe(__args: tuple[*_P], __fn1: Callable[[*_P], tuple[*_Q]], __fn2: Callable[[*_Q], _B]) -> _B:
+    ...
+
+
+@overload
+def starpipe(
+    __args: tuple[*_P],
+    __fn1: Callable[[*_P], tuple[*_Q]],
+    __fn2: Callable[[*_Q], tuple[*_X]],
+    __fn3: Callable[[*_X], _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], _B],
+) -> _B:
+    ...
+
+
+def starpipe(__args: Any, *__fns: Callable[..., Any]) -> Any:
     """Functional pipe_n (`||>`, `||>`, `|||>`, etc).
 
     Allows the use of function arguments on the left side of the
     function. Calls the function with tuple arguments unpacked.
 
     Example:
-        >>> starpipe((x, y), __fn) == __fn(x, y)  # Same as (x, y) ||> __fn
-        >>> starpipe((x, y), __fn, gn) == gn(fn(x))  # Same as (x, y) ||> __fn |> gn
+        >>> starpipe((x, y), fn) == fn(x, y)  # Same as (x, y) ||> fn
+        >>> starpipe((x, y), fn, gn) == gn(*fn(x))  # Same as (x, y) ||> fn |||> gn
+        >>> starpipe((x, y), fn, gn, hn) == hn(*gn(*fn(x)))  # Same as (x, y) ||> fn |||> gn ||> hn
         ...
     """
-    fn = fns[0] if len(fns) else starid
+    # Cast since unpacked arguments be used with TypeVarTuple
+    _starid = cast(Callable[..., Any], starid)
+    fn = __fns[0] if len(__fns) else _starid
 
-    return compose(*fns[1:])(fn(*args))
+    return starcompose(*__fns[1:])(fn(*__args))
 
 
 class PipeMixin:

tests/test_option.py

@@ -128,8 +128,6 @@ def test_option_order_none_none_works():
     assert not (xs < ys)
 
 
-
-
 def test_option_none_default_value():
     xs = Nothing
 
@@ -237,6 +235,27 @@ def test_option_some_map2_piped(x: int, y: int):
         case _:
             assert False
 
+def test_option_starmap_fluent():
+    xs = Some((42, 43))
+    mapper: Callable[[int, int], int] = lambda x, y: x + y
+    ys = xs.starmap(mapper)
+
+    match ys:
+        case Option(tag="some", some=value):
+            assert value == 85
+        case _:
+            assert False
+
+def test_option_starmap_piped():
+    xs = Some((42, 43))
+    mapper: Callable[[int, int], int] = lambda x, y: x + y
+    ys = pipe(xs, option.starmap(mapper))
+
+    match ys:
+        case Option(tag="some", some=value):
+            assert value == 85
+        case _:
+            assert False
 
 def test_option_some_bind_fluent():
     xs = Some(42)

tests/test_pipe.py

@@ -1,10 +1,15 @@
 from collections.abc import Callable
+from typing import TypeVar
 
 from hypothesis import given  # type: ignore
 from hypothesis import strategies as st
 
 from expression import pipe, pipe2
+from expression.core.pipe import starpipe, starid
 
+_A = TypeVar("_A")
+_B = TypeVar("_B")
+_C = TypeVar("_C")
 
 @given(st.integers())
 def test_pipe_id(x: int):
@@ -46,3 +51,30 @@ def test_pipe2_fn_gn(x: int, y: int):
     value = pipe2((x, y), fn, gn)
 
     assert value == gn(fn(x)(y))
+
+def test_starid_simple():
+    assert starid(1) == (1,)
+    assert starid(1, 2) == (1, 2)
+    assert starid(1, 2, 3) == (1, 2, 3)
+    assert starid(1, 2, 3, 4) == (1, 2, 3, 4)
+
+def fn(a: _A, b: _B) -> tuple[_A, _B]:
+    return a, b
+
+def gn(a: _A, b: _B) -> tuple[_B, _A]:
+    return b, a
+
+def yn(a: _A, b: _B) -> tuple[_A, _B, int]:
+    return a, b, 3
+
+def test_starpipe_simple():
+    assert starpipe((1, 2), fn) == fn(1, 2)
+
+def test_starpipe_id():
+    assert starpipe((1, 2), starid) == (1, 2)
+
+def test_starpipe_fn_gn():
+    assert starpipe((1, 2), fn, gn) == gn(*fn(1, 2))
+
+def test_starpipe_fn_gn_yn():
+    assert starpipe((1, 2), fn, gn, yn) == yn(*gn(*fn(1, 2)))

tests/test_result.py

@@ -424,7 +424,10 @@ def test_error_from_dict_works():
 def test_model_to_json_works():
     model = Model(one=Ok(10))
     obj = model.model_dump_json()
-    assert obj == '{"one":{"tag":"ok","ok":10},"two":{"tag":"error","error":{"message":"error"}},"three":{"tag":"error","error":{"message":"error"}}}'
+    assert (
+        obj
+        == '{"one":{"tag":"ok","ok":10},"two":{"tag":"error","error":{"message":"error"}},"three":{"tag":"error","error":{"message":"error"}}}'
+    )
 
 
 def test_error_default_value():

tests/test_seq.py

@@ -382,3 +382,4 @@ def test_seq_monad_law_associativity_empty(value: int):
     # Empty list
     m = empty
     assert list(m.collect(f).collect(g)) == list(m.collect(lambda x: f(x).collect(g)))
+