Annotations for variable length function composition/pipeline.

[randolf-scholz]

class Transform(Protocol[X, Y]):
    def __call__(self, x: X, /) -> Y: ...

class Pipeline(Transform):
    funcs: list[Transform]  # <-- how to annotate?
    
    def __init__(self, funcs: Iterable[Transform], /) -> None:  # <-- how to annotate?
        self.funcs = list(funcs)

    def __call__(self, x):
        for func in self.funcs:
            x = func(x)
        return x

Is there any simple way to have more precise type hints for Pipeline, without going completely overboard with overloads?
In particular, it would be nice to have:

1. Pipeline(Transform[X_first, Y_last])
2. some type checking that the domain/codomain of each transform is compatible
3. Handling of the edge cases with 0 or 1 arguments: Pipeline([]) and Pipeline([Transform[X,Y]]).

The best I could come up with handles (1) and (3), but requires changing the __init__-signature from (self, funcs: Iterable[Transform], /) to self, *funcs: *tuple[Transform, ...], /), and the introduction of an additional subclass IdentityPipe.

I wonder if there is any simpler/better way of doing it.

---

Code sample in pyright playground

from typing import Protocol, TypeVar, overload, Any, cast, reveal_type, Never, assert_type

X = TypeVar("X")
Y = TypeVar("Y")
U = TypeVar("U")
V = TypeVar("V")

class Transform(Protocol[X, Y]):
    def __call__(self, x: X, /) -> Y: ...

class Pipeline(Transform[U, V]):
    @overload
    def __new__(cls, *funcs: *tuple[()]) -> "IdentityPipe": ...
    @overload
    def __new__(cls, *funcs: *tuple[Transform[X, Y]]) -> "Pipeline[X, Y]": ...
    @overload
    def __new__(cls, *funcs: *tuple[Transform[X, Any], *tuple[Transform, ...], Transform[Any, Y]]) -> "Pipeline[X, Y]": ...
    def __new__(cls, *funcs: *tuple[Transform, ...]):  # pyright: ignore[reportInconsistentOverload]
        return super().__new__(cls, *funcs)

    
    def __init__(self, *funcs: *tuple[Transform, ...]) -> None:  # <-- how to annotate?
        self.funcs = list(funcs)

    def __call__(self, x: U, /) -> V:
        x = cast(Any, x)
        for func in self.funcs:
            x = func(x)
        return cast(V, x)

class IdentityPipe(Pipeline[Any, Any]):
    funcs: list[Never]
    def __call__(self, x: X) -> X:
        return x

S: Transform[int, float] = lambda x: float(x)
Z: Transform[float, float] = lambda x: 2*x
T: Transform[float, str] = lambda x: str(x)

x: int = 1

P0 = Pipeline()
assert_type(P0, IdentityPipe)
assert_type(P0(x), int)

P1 = Pipeline(S)
assert_type(P1, Pipeline[int, float])
assert_type(P1(x), float)

P2 = Pipeline(S, T)
assert_type(P2, Pipeline[int, str])
assert_type(P2(x), str)

P3 = Pipeline(S, Z, T)
assert_type(P3, Pipeline[int, str])
assert_type(P3(x), str)

[decorator-factory]

Right now that is not possible. However, you can chain method calls to achieve the same result:

from typing import Generic, TypeVar
from collections.abc import Callable

A = TypeVar("A")
B = TypeVar("B")
C = TypeVar("C")

class Pipe(Generic[A, B]):
    def __init__(self, fn: Callable[[A], B]) -> None:
        self._fn = fn

    def __call__(self, arg: A, /) -> B:
        return self._fn(arg)
    
    def then(self, fn: Callable[[B], C], /) -> "Pipe[A, C]":
        return Pipe(lambda a: fn(self._fn(a)))

Then you can do pipe = Pipe(S).then(Z).then(T) and have pipe inferred as Pipe[int, str]

[randolf-scholz]

I do something like this, in my code the Transform-derived base class has a method for chaining that returns Pipeline(self, other)