TypeVarTuple and Cartesian Product of functions

[randolf-scholz]

Consider the Cartesian product of functions

def prod_fn(*funcs):
    r"""Cartesian Product of Functions.

    It is assumed every function takes a single positional argument.
    """

    def __prod_fn(args, /):
        """Argument is a tuple with the input for each function."""
        return tuple(f(arg) for f, arg in zip(funcs, args))

    return __prod_fn

How to type hint it? It should be something along the lines of

def prod_fn(*funcs: *tuple[???]) -> Callable[[tuple[*U]], tuple[*V]]:
    r"""Cartesian Product of Functions.

    It is assumed every function takes a single positional argument.
    """

    def __prod_fn(args: tuple[*U], /) -> tuple[*V]:
        """Argument is a tuple with the input for each function."""
        return tuple(f(arg) for f, arg in zip(funcs, args))

    return __prod_fn

Where U and V are TypeVarTuples for the inputs and outputs of each individual function. But I fail to see how to express the concept that funcs is the variadic tuple tuple[Callable[[U₁], V₁]], Callable[[U₂], V₂]], ..., Callable[[Uₙ], Vₙ]]]

[Daverball]

This just isn't possible yet, this would require something like a Map construct to apply a type to each element of the TypeVarTuple, (although this case is even more complex since there's two TypeVarTuple that both need to be mapped at the same time)

Right now you can only unpack a TypeVarTuple directly into other variadic generics (i.e. tuple, Callable, Union and any user-defined generics that take a TypeVarTuple).

P.S. Wouldn't you need to use varargs for it to actually be the cartesian product of functions? The definition to me looks like the resulting function is supposed to have as many arguments as input functions, not one argument that's a tuple of all the inputs.

def prod_fn(*funcs: *Map[(U, V), Callable[[U], V]]]) -> Callable[[*U], tuple[*V]]:
    r"""Cartesian Product of Functions.

    It is assumed every function takes a single positional argument.
    """

    def __prod_fn(*args: *U) -> tuple[*V]:
        """Argument is a tuple with the input for each function."""
        return tuple(f(arg) for f, arg in zip(funcs, args))

    return __prod_fn

[Daverball]

In the meantime you can use an approach like itertools.product and use overloads to cover the most common cases:

@overload
def prod_fn(f1: Callable[[T1], R1]]], /) -> Callable[[T1], tuple[R1]]: ...  # 1 Callable
@overload
def prod_fn(f1: Callable[[T1], R1]]], f2: Callable[[T1], R2], /) -> Callable[[T1, T2], tuple[R1, R2]]: ...  # 2 Callables
@overload
def prod_fn(f1: Callable[[T1], R1]]], f2: Callable[[T2], R2], f3: Callable[[T3], R3], /) -> Callable[[T1, T2, T3], tuple[R1, R2, R3]]: ...  # 3 Callables

# Maybe add a few more for 4-6 Callables
# To avoid overlap with the first n overloads and increase robustness
# you could first add n positional arguments, matching the `*args` annotation to the following two overloads
# although this may cause issues when unpacking an arbitrary length iterable i.e. you call it like `prod_fn(*[f1, f2, f3, ...])`

@overload
def prod_fn(*funcs: Callable[[T], R]]]) -> Callable[[*tuple[T, ...]], tuple[R, ...]]: ...  # arbitrary length, uniform types
@overload
def prod_fn(*funcs: Callable[[Any], Any]]]) -> Callable[..., tuple[Any, ...]]: ...  # catch-all

def prod_fn(*funcs: Callable[[Any], Any]]]) -> Callable[..., tuple[Any, ...]]:
    r"""Cartesian Product of Functions.

    It is assumed every function takes a single positional argument.
    """

    def __prod_fn(*args: Any) -> tuple[Any, ...]:
        """Argument is a tuple with the input for each function."""
        return tuple(f(arg) for f, arg in zip(funcs, args))

    return __prod_fn