Does mypy have a Subclass-Acceptable Return Type?

Question

I'm wondering how (or if it is currently possible) to express that a function will return a subclass of a particular class that is acceptable to mypy?

Here's a simple example where a base class Foo is inherited by Bar and Baz and there's a convenience function create() that will return a subclass of Foo (either Bar or Baz) depending upon the specified argument:

class Foo:
    pass


class Bar(Foo):
    pass


class Baz(Foo):
    pass


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


bar: Bar = create('bar')

When checking this code with mypy the following error is returned:

error: Incompatible types in assignment (expression has type "Foo", variable has type "Bar")

Is there a way to indicate that this should be acceptable/allowable. That the expected return of the create() function is not (or may not be) an instance of Foo but instead a subclass of it?

I was looking for something like:

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

but that doesn't exist. Obviously, in this simple case, I could do:

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

but I'm looking for something that generalizes to N possible subclasses, where N is a number larger than I want to be defining as a typing.Union[...] type.

Anyone have any ideas on how to do that in a non-convoluted way?

---

In the likely case that there is no non-convoluted way to do this, I am aware of a number of less-than-ideal ways to circumvent the problem:

1. Generalize the return type:

def create(kind: str) -> typing.Any:
    ...

This resolves the typing problem with the assignment, but is a bummer because it reduces the type information of the function signature's return.

2. Ignore the error:

bar: Bar = create('bar')  # type: ignore

This suppresses the mypy error but that's not ideal either. I do like that it makes it more explicit that the bar: Bar = ... was intentional and not just a coding mistake, but suppressing the error is still less than ideal.

3. Cast the type:

bar: Bar = typing.cast(Bar, create('bar'))

Like the previous case, the positive side of this one is that it makes the Foo return to Bar assignment more intentionally explicit. This is probably the best alternative if there's no way to do what I was asking above. I think part of my aversion to using it is the clunkiness (both in usage and readability) as a wrapped function. Might just be the reality since type casting isn't part of the language - e.g create('bar') as Bar, or create('bar') astype Bar, or something along those lines.

Answer 1

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')

Answer 2

I solved a similar issue using typing.Type. For your case I would use it as so:

class Foo:
    pass

class Bar(Foo):
    pass

class Baz(Foo):
    pass

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

It seems this works because Type is "covariant" and though I'm no expert the above link points to PEP484 for more details

Answer 3

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.