CODEWITHSUNDEEP
pythonpython-typing
Stoney
Stoney

Reputation: 407

How to infer the type of last element passed to `*args` parameter in python?

I want to write a pipeline function which takes an abritrary number of functions and forwards the input of the first to the second and so on up until the last. The output of pipeline is a Callable. The input of that Callable matches the first function given to pipeline. The output matches the last function given to pipeline. I would like to add type annotations to pipeline. This is what I came up with so far:

from typing import TypeVar
from typing_extensions import ParamSpec, Concatenate

P = ParamSpec("P")
Output = TypeVar("Output")

def pipeline(func: Callable[Concatenate[P], T], *funcs: Callable[..., Output]) -> Callable[Concatenate[P], Output]:
    ...

But this gives a Union of all Output types of *funcs. Example:

def double(x: int) -> Tuple[int, int]:
    return x, x*2

def add(x: int, y: int) -> int:
    return x + y

def to_string(x: int) -> str:
    return str(x)

new_func = pipeline(add, double, add, to_string)

With the type annotation shown above the type / signature of new_func results in

new_func: (x: int, y: int) -> (Tuple[int, int] | int | str)

It's quite clear that pipeline should have the following signature:

new_func: (x: int, y: int) -> str

Is there a way to achieve this with type annotations?

Upvotes: 1

Views: 256

Answers (2)

ruohola
ruohola

Reputation: 24155

I'm quite sure Python 3.10's type annotations are not powerful enough so you could define pipeline to support an arbitrary number of arguments. But, you could do it with overloads (@typing.overload) and define pipeline's signature separately for each arity.

You could take reference from the popular TypeScript library fp-ts, and more specifically from its flow function: https://github.com/gcanti/fp-ts/blob/9da2137efb295b82fb59b7b0c2114f2ceb40a2b5/src/function.ts#L228-L351

Upvotes: 0

Stoney
Stoney

Reputation: 407

Actually found a way to solve this problem using TypeVarTuple and Unpack. PEP646 provides an example that shows how to unpack heterogenous parameters passed as *args. With this approach the example from the question could be typed as follows:

from typing import TypeVar
from typing_extension import TypeVarTuple, ParamSpec, Unpack

Output = TypeVar("Output")
InputParams = ParamSpec("InputParams") 
IntermediateFunctions = TypeVarTuple("IntermediateFunctions")

def pipeline(*functions: Unpack[
  Tuple[
    Callable[InputParams, Any], 
    Unpack[IntermediateFunctions], 
    Callable[..., Output]
  ]
]) -> Callable[InputParams, Output]:
  ...

# From python 3.11 onwards you can use * instead of Unpack
def pipeline(*functions: *Tuple[Callable[InputParams, Any], *IntermediateFunctions, Callable[..., Output]]) -> Callable[InputParams, Output]:
  ...

Running mypy against this implementation failed as Unpack is not supported yet. I described the issue here

Upvotes: 2

Related Questions