Should mypy infer type from Union options?

Question

In the following code fn_int will only be called with an int as an argument, yet mypy complains I might be passing it a str (viceversa for fn_str).

Is there something I'm missing here? Should/can I somehow narrow down the type before passing arguments to fn_int and fn_str?

from typing import Union


def fn_int(arg: int) -> None:
    print(arg)


def fn_str(arg: str) -> None:
    print(arg)


def fn(arg: Union[str, int]) -> None:

    if arg in range(4):
        fn_int(arg)
    elif arg == "str":
        fn_str(arg)
    else:
        raise ValueError

> mypy mypy_test.py
mypy_test.py:15: error: Argument 1 to "fn_int" has incompatible type "Union[str, int]"; expected "int"  [arg-type]
mypy_test.py:17: error: Argument 1 to "fn_str" has incompatible type "Union[str, int]"; expected "str"  [arg-type]

Answer 1

What you're looking for is called type narrowing, and mypy doesn't do type narrowing for == or in comparisons. After all, they don't really guarantee type. As a fairly common example, after x == 3 or x in range(5), x could still be a float rather than an int. (Sure, a standard string won't pass an in range(4) check, and a standard int won't pass an == "str" check, but you could have a weird subclass.)

mypy will do type narrowing for the following kinds of expressions:

• isinstance(x, SomeClass) narrows x to SomeClass.
• issubclass(x, SomeClass) narrows x to Type[SomeClass].
• type(x) is SomeClass narrows x to SomeClass.
• callable(x) narrows x to callable type.

You can also write your own type guards with typing.TypeGuard, and mypy will understand those too.

Answer 2

arg is str or int, but you never check which type it is, therefore check with if isinstance(arg, int).

then mypy knows that this function is only called if it is of desired type.

Example:

if isinstance(arg, int):
    fn_int(arg)
elif isinstance(arg, str):
    fn_str(arg)
else:
    raise ValueError