Python function with two overloads calls another one with the same overloads - explain to mypy

Question

How do we explain to mypy that this code is ok?

from typing import overload, Union


@overload
def foo(x: float, y: str):
    ...

@overload
def foo(x: list, y: tuple):
    ...

def foo(x: Union[float, list], y: Union[str, tuple]):
    bar(x, y)  # type error!


@overload
def bar(x: float, y: str):
    ...

@overload
def bar(x: list, y: tuple):
    ...

def bar(x: Union[float, list], y: Union[str, tuple]):
    print(x)
    print(y)

The type errors are:

foo.py:13: error: Argument 1 to "bar" has incompatible type "Union[float, List[Any]]"; expected "float"  [arg-type]
foo.py:13: error: Argument 2 to "bar" has incompatible type "Union[str, Tuple[Any, ...]]"; expected "str"  [arg-type]

In similar cases I would use TypeVar, but the problem is that here x and y are different (but correlated) types.

Answer 1

Solution 1: Combine related types

If it is possible to change foo's signature, then we can describe the type relationships to mypy as follows:

@overload
def foo(xy: Tuple[float, str]):
    ... 

@overload
def foo(xy: Tuple[list, tuple]):
    ... 

def foo(xy: Union[Tuple[float, str], Tuple[list, tuple]]):
    bar(*xy)

# bar and its overloads unchanged

Solution 2: Secretly combine related types

If we can't change foo's signature, then as far as I can tell, this is an example of a situation that cannot at present be completely captured by mypy's type system¹. Still, we can keep the amount of code excepted from type-checking to a single line

@overload
def foo(x: float, y: str):
    ... 

@overload
def foo(x: list, y: tuple):
    ... 

def foo(x, y):        # no type annotations, so...
    foo_impl((x, y))  # ...this line is not type-checked

def foo_impl(xy: Union[Tuple[float, str], Tuple[list, tuple]]):
    bar(*xy)          # passes type-checking

# bar and its overloads unchanged

If foo is a single line, then we can achieve a similar result by simply removing type annotations on the signature of the implementation of foo

@overload
def foo(x: float, y: str):
    ... 

@overload
def foo(x: list, y: tuple):
    ... 

def foo(x, y):  # no type annotations, so...
    bar(x, y)   # ...this line is not type-checked

Keeping information about type correlations

Intuitively, the initial type annotations of the arguments of foo can be thought of as jointly specifying the type

Tuple[Union[float, list], Union[str, tuple]]

which discards the information about the correlation between the type of x and the type of y. We can keep that information by "transposing" Tuple and Union, i.e. by saying

Union[Tuple[float, str], Tuple[list, tuple]]

which preserves the information about the correlation. Both solutions above achieve this preservation in slightly different ways.

---

¹ There do exist type systems sophisticated enough to describe the above situation. For example, in C++ one can describe correlations between types using template specialization.