Search code examples
pythonmypy

Mypy type narrowing from assertion


I have two variables a and b that are either int or str. I write an assertion that insists a and b are either both ints or strs.

If I typenarrow a to an int, is there a way for mypy to infer b is also an int? Here is some sample code.

Mypy version:

mypy 0.980+dev.0f17aff06ac1c05c442ba989e23655a2c6adbfbf (compiled: no)

Thanks for your help.

def my_func(a: int | str, b: int | str):
    # We assert one of the statements is true: 1) a and b are ints, or 2) a and b are strings.
    # In the case of an int a, and a string b, or vice versa, this assertion will fail.
    assert isinstance(a, int) == isinstance(b, int)
    # Another alternative assertion
    assert type(a) == type(b)

    if isinstance(a, int):
        reveal_type(b)  # mypy still thinks b is int or str


Solution

  • Using typing.TypeVar:

    from typing import TypeVar
    
    T = TypeVar('T', int, str)
    
    
    def reveal_type(a: int):
        pass
    
    
    def foo(a: str):
        pass
    
    
    def my_func(a: T, b: T):
        if isinstance(a, int):
            reveal_type(b)  # pass
        else:
            foo(b)          # pass
    

    If we simply exchange the calling positions of the two functions, mypy will find that they are all wrong calls and give two errors:

    def my_func(a: T, b: T):
        if isinstance(a, int):
            foo(b)          # Argument 1 to "foo" has incompatible type "int"; expected "str" (16:12)
        else:
            reveal_type(b)  # Argument 1 to "reveal_type" has incompatible type "str"; expected "int" (18:20)