[Type hint] Avoid massive `@overlad`ing for multiple `Union`-typed inputs

Hello everyone.

Imagine I define such function,

def f(
    x: tuple[X, ...],
    y: tuple[Y, ...],
    z: tuple[Z, ...],
) -> tuple[tuple[tuple[T, ...], ...], ...]: ...

where f applies some logic for every combination of elements from x, y and z to return a T element.

Now, I want to extend the API of f by allowing, for each argument (e.g. x), passing a single bare element – such as x: tuple[X, ...] | X --, while the output will be squeezed accordingly to reflect this.

For example with X = Y = Z = T = int, f(1, 2, 3) would return a simple int instead of a tuple[tuple[tuple[int, ...], ...], ...].

There are basically 8 possible signatures, that I can directly define with @overload.

@overload
def f(x: X, y: Y, z: Z) -> T: ...

@overload
def f(x: tuple[X, ...], y: Y, z: Z) -> tuple[T, ...]: ...

# etc... 8 times

Is it possible to find a better way to “dynamically” emulate this overloading possibilities? Because in addition to be very ugly, it also becomes impractical very fast.

Thanks in advance!

In general, this is not currently possible. There are some special cases where signatures similar to this can be written using things like constrained type vars, but those wouldn’t help with this particular function. Functions where the “shape” of input/output containers is dynamic unfortunately cannot be typed very well right now. There is work being done on this currently, but from what I’ve understood about that even the proposed solutions there wouldn’t help with this particular example (though they would solve many very similar problems).

1 Like