Should mypy infer type from Union options?

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]

>Solution :

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.