CODEWITHSUNDEEP
pythonmypy
hwong557
hwong557

Reputation: 1439

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

Upvotes: 1

Views: 652

Answers (1)

Mechanic Pig
Mechanic Pig

Reputation: 7736

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)

Upvotes: 3

Related Questions