Python typing: Narrowing type from function that returns a Union

I have difficulties finding return types that satisfy mypy. I have two functions. The first one returns a Union type, since the type depends on the parameter given to the function. The second function calls the first one with a default parameter. Thus, the type is not a Union type -- it can be narrowed down to one of the types in the Union.

Let me give you an example of a very simple version of my problem:

from typing import Union


def first(as_string: bool) -> Union[str, int]:
    if as_string:
        return "42"
    else:
        return 42


def second(as_string: bool) -> str:
    return first(True)

This causes an error:

Incompatible return value type (got "str", expected "Union[str, int]")

How can I prevent mypy from throwing an error while still using type hints?

If you want to suggest splitting the first function, please take in mind that this is just a simplification. My (first) function receives a function (sklearn.metrics function) and will return a scalar most of the time. Only when applying a confusion matrix, the type changes. The first function does some preprocessing and than applies the metric. I simply want to have a differently named function for the confusion matrix since I think it is a special case of a metric.

Solution

mypy can't infer the relationship between the types of the parameters and the types of the return values.

You have two options:

Use an assert to make sure that the type is correct:

def second(as_string) -> str:
    ret = first(True)
    assert isinstance(ret, str)
    return ret

You can use typing.cast to assert to the typechecker without altering runtime behavior (i.e. without introducing the potential for an AssertionError to be raised).
```
from typing import cast

def second(as_string) -> str:
    return cast(str, first(True))
```