CODEWITHSUNDEEP
pythonunionsubclasstyping
waszil
waszil

Reputation: 470

Python how to define type hint for union of all subclasses

Let's say I have a class that have multiple subclasses in a recursive/tree hierarchy:

class Animal:
    pass
class Cat(Animal):
    pass
class HouseCat(Cat):
    pass

And I have a function, that can create instances of these classes based on some condition:

from typing import Union
def creator(condition) -> Union[Animal, Cat, HouseCat]:
    # for codition1
    return Animal()
    # for codition2
    return Cat()
    # ...etc...

My problem is that PyCharm shows a warning for something like this, if I only use -> Animal: as return value annotation:

my_list: List[Cat] = []
obj = creator(...)
my_list.append(obj)  # <-- Expected type 'Cat' (matched generic type '_T'), got 'Animal' instead

Is there a way to define a type hint instead of manually writing a Union for all subclasses?

Upvotes: 3

Views: 2271

Answers (1)

abhijat
abhijat

Reputation: 605

You can use the TypeVar construct to write code with bounds:

from typing import TypeVar, List


class Animal:
    pass


class Cat(Animal):
    pass


class HouseCat(Cat):
    pass


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


def creator() -> A:
    pass


my_list: List[Cat] = []
obj = creator()
my_list.append(obj)

https://docs.python.org/3/library/typing.html#typing.TypeVar

Upvotes: 2

Related Questions