Search code examples
pythonpython-typingmypy

Python typing: How to assert which type from a union of types a value has?


In short, I have a function that returns either int or float. The caller function then checks the return-type of the first function and return -1 if float else return the original value since it must be int.

# pseudo code for the aforementioned 

def f1(*args, **kwargs) -> int | float: ... 

def f2(*args, **kwargs) -> int:
    ans = f1(...)
    if type(ans) == float:
        return -1
    return ans  # the idea here is that if f1 does not return float, it must return an int which is a valid return for f2

My static checker fails with the following error

Expression of type "int | float" cannot be assigned to return type "int"
  Type "int | float" cannot be assigned to type "int"
    "float" is incompatible with "int"

The error message is very straight forwards, f1 returns either int or float, since f2 expects to return int, it can't return the result of f1 directly. However, (ideally) my if statement guards against the possibility that the result of f1 is float.

Does anyone know of a better way to annotate the following. I am currently using the type: ignore flag but I wish not to use this workaround.


Solution

  • You need to check with isinstance:

    def f1(*args, **kwargs) -> int | float: ... 
    
    def f2(*args, **kwargs) -> int:
        ans = f1(...)
        if isinstance(ans, float):
            return -1
        # now the typechecker can infer the type of 'ans' as int
        return ans 
    

    More info in Mypy documentation