Narrowing down a type from function return types

Question

I have a function which might return values of two different types, e.g. Union[str, int]. The type of the returned value can be determined from the function argument. So on each call, I know which type should be returned. However, mypy complains that Union[str, int] type cannot be assigned to a stricter type, e.g. str. Is there a way to allow stricter type assignment, without completely turning type checking off (e.g. with # type: ignore comment)?

For example, if I have this function:

from typing import Union


def change_value(action: str, value: int) -> Union[str, int]
    if action == "stringify":
        return str(value)
    return value

stringified_value: str
stringified_value = change_value("stringify", 10)

Mypy complains with this error:

Expression of type "str | int" cannot be assigned to declared type "str"

I would like to remove mypy error by somehow telling mypy, that in this case type narrowing down is always correct.

Answer 1

It can be done using overloads:

from typing import Literal, Union, overload

@overload
def change_value(action: Literal['stringify'], value: int) -> str: ...
@overload
def change_value(action: Literal['whatever_else'], value: int) -> int: ...

def change_value(action: str, value: int) -> Union[str, int]:
    if action == "stringify":
        return str(value)
    return value

stringified_value: str = change_value("stringify", 10)

playground

This approach has an additional benefit: you explicitly tell that stringify and whatever_else are two only possible actions, and call like change_value('stingify', 2) will be denied by type checker (noticed the typo? If not, it would be hard to debug otherwise).