Dependent types and polymorphism in Python with mypy

For the following example, mypy returns an error:

error: Incompatible types in assignment (expression has type "A", variable has type "A1")

from typing import Type

class A:
    pass

class A1(A):
    pass

class A2(A):
    pass

def fun(A_type: Type[A]) -> A:
    if A_type == A1:
        return A1()
    else:
        return A2()

a1: A1 = fun(A1)

What I would ideally like to do is to enforce a dependency in the signature of fun:

def fun(A_type: Type[A]) -> A_type

Is this possible; if not, what is recommended (note: I want this to work for as yet undefined sub-classes of A, so I don't think I can use the overload decorator)? Is my best option just to use cast?

Use a TypeVar with a bound on it:

https://mypy.readthedocs.io/en/latest/generics.html#type-variables-with-upper-bounds

from typing import Type, TypeVar

class A:
    pass

class A1(A):
    pass

class A2(A):
    pass

T_A = TypeVar('T_A', bound='A')

def fun(A_type: Type[T_A]) -> T_A:
    if A_type == A1:
        r1 = A1()
        assert isinstance(r1, A_type)
        return r1
    else:
        r2 = A2()
        assert isinstance(r2, A_type)
        return r2

a1: A1 = fun(A1)
a2: A2 = fun(A2)
print("winner winner chicken dinner")

typechecks clean and runs without failing either type assert:

C:\test\python>mypy polymorph.py
Success: no issues found in 1 source file

C:\test\python>python polymorph.py
winner winner chicken dinner

In this example the type T_A is required to be a subclass of A, but it's a particular type, and the typing of fun requires that it returns the same type it receives as an argument.

Unfortunately the static type checker isn't quite smart enough to bind the type unless you add the runtime assert in there (there might be some way to do this better with the Generic type but it eludes me).