Mypy Assertions on the Types of Types

Viewed 924

In Python, we can assign a type to an variable in this way and pass mypy's typecheck:

from typing import Type, Union


class Foo:
    pass


MyType: Type[Foo] = Foo

Similarly, we can also use Union for typing

from typing import Type, Union


class Foo:
    pass


class Bar:
    pass


def func(is_bar: bool) -> Union[Type[Foo], Type[Bar]]:
    if is_bar:
        return Bar
    else:
        return Foo

Knowing the value of is_bar, it would be totally reasonable to use assertions after calling func for type narrowing.

However, assertions do not seem to work at all for MyType as mypy is giving an error for the following code example:

MyType = func(True)

assert MyType is Bar  
# `assert MyType is Type[Bar]` gives the same result

Test: Bar = MyType  # Incompatible types in assignment (expression
                    # has type "Union[Type[Foo], Type[Bar]]",
                    # variable has type "Bar")

cast and isinstance do not work either.

MyType = func(True)
cast(Type[Bar], MyType)
Test: Bar = MyType  # MyType is considered `Type[Bar] | Type[Foo]`
MyType = func(True)
assert isintance(MyType, Type[Bar])
Test: Bar = MyType  # MyType is considered `Type[Bar] | Type[Foo]`

My questions is: how to do type narrowing for the types of types? Or is it a limitation of mypy? If that's so, what would be the workaround?

Related:

2 Answers

The cast helper returns its constrained argument, it does not actually change its argument to be constrained.

cast to the desired Type[...] and assign or use the result:

Test = cast(Type[Bar], MyType)
reveal_type(Test)  # note: Revealed type is "Type[so_testbed.Bar]"

Another solution, that avoids having to use typing.cast or isinstance, is to use typing.overload, which allows you to register multiple signatures of a single function. All functions decorated with @typing.overload are ignored at runtime in favour of the "concrete" implementation, so the body of these functions can just be a literal ellipsis. By combining typing.overload with typing.Literal, we can register one signature of the function if the value True is passed in, and another if the value False is passed in:

from typing import Type, Union, overload, Literal


class Foo:
    pass


class Bar:
    pass


@overload
def func(is_bar: Literal[True]) -> Type[Bar]: ...


@overload
def func(is_bar: Literal[False]) -> Type[Foo]: ...


def func(is_bar: bool) -> Union[Type[Foo], Type[Bar]]:
    if is_bar:
        return Bar
    else:
        return Foo
        

test: Type[Bar] = func(True)
test2: Type[Foo] = func(False)
test3: Bar = func(True)()
test4: Foo = func(False)()

It type-checks fine.

Related