Does mypy have a Subclass-Acceptable Return Type?

You can find the answer here. Essentially, you need to do:

class Foo:
    pass


class Bar(Foo):
    pass


class Baz(Foo):
    pass

from typing import TypeVar
U = TypeVar('U', bound=Foo)

def create(kind: str) -> U:
    choices = {'bar': Bar, 'baz': Baz}
    return choices[kind]()


bar: Bar = create('bar')

Mypy is not complaining about the way you defined your function: that part is actually completely fine and error-free.

Rather, it's complaining about the way you're calling your function in the variable assignment you have at your very last line:

bar: Bar = create('bar')

Since create(...) is annotated to return a Foo or any subclass of foo, assigning it to a variable of type Bar is not guaranteed to be safe. Your options here are to either remove the annotation (and accept that bar will be of type Foo), directly cast the output of your function to Bar, or redesign your code altogether to avoid this problem.


If you want mypy to understand that create will return specifically a Bar when you pass in the string "bar", you can sort of hack this together by combining overloads and Literal types. E.g. you could do something like this:

from typing import overload
from typing_extensions import Literal   # You need to pip-install this package

class Foo: pass
class Bar(Foo): pass
class Baz(Foo): pass

@overload
def create(kind: Literal["bar"]) -> Bar: ...
@overload
def create(kind: Literal["baz"]) -> Baz: ...
def create(kind: str) -> Foo:
    choices = {'bar': Bar, 'baz': Baz}
    return choices[kind]()

But personally, I would be cautious about over-using this pattern -- I view frequent use of these types of type shenanigans as something of a code smell, honestly. This solution also does not support special-casing an arbitrary number of subtypes: you have to create an overload variant for each one, which can get pretty bulky and verbose.