Overloads with variable-length arguments via *args

[BaconPancakes]

I'm trying to specify overloads for a function that returns a tuple of the same length as the number of arguments given, while also supporting a variable length tuple return for a *args invocation. Is there any way to make this work?

Code sample in pyright playground

from typing import overload


@overload
def my_fun(bar1: int, /) -> int: # type error
    ...

@overload
def my_fun(bar1: int, bar2: int, /) -> tuple[int, int]:
    ...

@overload
def my_fun(*bar: int) -> tuple[int, ...]:
    ...

def my_fun(*bar: int,) -> tuple[int, ...] | int:
    ...

my_iter = (1,2,3)
result = my_fun(*my_iter)

Note that in my actual code, the function is also generic on argument and return, so with the lack of higher-kind TypeVar I'm not able to use some kind of TypeVarTuple solution either.

[erictraut]

If I understand your problem correctly, if you pass your function a single int argument, it should return an int. If you pass it two or more int arguments, it should return a tuple that contains all of the values passed to it. If that's correct, then the following should address your needs:

@overload
def my_fun(a: int, /) -> int:  ...
@overload
def my_fun(a: int, b: int, /, *bar: int) -> tuple[int, int, *tuple[int, ...]]: ...

def my_fun(*bar: int) -> tuple[int, ...] | int: ...

[BaconPancakes]

Your interpretation is correct given my example, but I was not clear enough. Let me expand on my example with the actual container objects that I'm using with TypeVars.

What I want to accomplish is a generic function that accepts parameters of distinct types, and returns a tuple of containers bound to those input types in the same order. I don't believe this is generally accomplishable, so what I have in my actual code is a workaround of 10 overloads.

In addition, I want to be able to pass this function an Iterator, such that it returns a variable length Tuple. This is where I got stuck with my above approach.

Code sample in pyright playground

from typing import overload, Any

class Container[T]:
    ...

@overload
def my_fun[T](bar1: T, /) -> Container[T]:
    ...

@overload
def my_fun[T1, T2](bar1: T1, bar2: T2, /) -> tuple[Container[T1], Container[T2]]:
    ...

@overload
def my_fun[T1, T2, T3](bar1: T1, bar2: T2,  bar3: T3,/) -> tuple[Container[T1], Container[T2], Container[T3]]:
    ...

# T1, T2, T3, T4, ... etc

def my_fun(*bar: Any) -> tuple[Container[Any], ...] | Container[Any]:
    ...

a, b, c = my_fun(1, False, "hello") # a: Container[int], b: Container[bool], c: Container[str]

my_iter = range(0,5)
result = my_fun(*my_iter) # returns Container[int] instead of tuple[Container[Any], ...]
# Is there a better option than typing.cast() ?

[erictraut]

You are correct that this isn't possible with the type system today. There has been some discussion of a "Map" mechanism that allows a transform to be applied to each element in a TypeVarTuple, but this hasn't made it to a PEP proposal yet.

[BaconPancakes]

Thanks for the expert advice. I meant to open this thread in pyright actually - thanks for responding regardless.