Narrowing argument types based on overload variants

[bushkov]

Feature

Currently mypy cannot narrow down argument types in a function implementation based on function's overload variants. The following code

@overload
def f(x: None, y: None) -> None:...

@overload
def f(x: str, y: int) -> int:...

def f(x: Optional[str], y: Optional[int]) -> Optional[int]:
    if x is not None:
        return len(x) + y # error: Unsupported operand types for + ("int" and "None")
    else:
        return None

results in a Unsupported operand types for + ("int" and "None")" error. Despite of the fact that f() can be called only with either (None, None) or (str, int) which is specified by the respective overloads, mypy cannot infer that if the first argument is not None, then the other argument also shouldn't be None.

Here is my use case where I find this inference would be helpful:

class A:
    pass

class B:
    pass

class C:
    pass

class D:
    pass

@overload
def f(x: A) -> C:...

@overload
def f(x: B) -> D:...

def f(x: Union[A, B]) -> Union[C, D]:
    if isinstance(x, A):
        return C()
    else:
        return D()

@overload
def g(x: A, func: Callable[[C], bool]) -> None:...

@overload
def g(x: B, func: Callable[[D], bool]) -> None:...

def g(x: Union[A, B], func: Union[Callable[[C], bool], Callable[[D], bool]]) -> None:
    obj = f(x)
    func(obj) # error: Argument 1 has incompatible type "Union[C, D]"; expected "C"

Without being able to infer that when x is an instance of A, func should be an instance of Callable[[C], bool] (which is specified by the overload variants), mypy results in the following error:

test.py:37: error: Argument 1 has incompatible type "Union[C, D]"; expected "C"
test.py:37: error: Argument 1 has incompatible type "Union[C, D]"; expected "D"