Python typing: Narrowing type from function that returns a Union

user_10652408 0
Viewed 799

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.

mypy typing python
1 Answers

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

You have two options:

  1. 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

  2. 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))
user_177698150
Related