Python type annotations with class attribute of type List of derived classes

Question:

For classes that contain as an attribute a list of instances of other derived classes, what is the correct way to type annotate.
Bellow is a "simple" example with two pairs of base and derived classes where I can reproduce the mypy errors I have troubles solving.

from dataclasses import dataclass
from typing import List, TypeVar

ClassType = TypeVar("ClassType", bound="BaseClass")


@dataclass
class BaseClass:
    a: int
    b: int


@dataclass
class DerivedClass(BaseClass):
    c: str


@dataclass
class BaseExample:
    list_attr: List[ClassType]  # line 20

    def validate(self, attrs: List[ClassType]) -> List[ClassType]:
        return [attr for attr in attrs if attr.a > 0]


@dataclass
class DerivedExample:
    list_attr: List[ClassType]  # Line 28  # Is actually List[DerivedClass]
    other_attr: int

    def select(self, attrs: List[ClassType]) -> List[ClassType]:
        return [attr for attr in attrs if attr.c] # line 32


ex = DerivedExample(list_attr=[DerivedClass(a=1, b=1, c="text")], other_attr=0)

I get the following mypy errors:

example.py:20: error: Type variable "example.ClassType" is unbound
example.py:20: note: (Hint: Use "Generic[ClassType]" or "Protocol[ClassType]" base class to bind "ClassType" inside a class)
example.py:20: note: (Hint: Use "ClassType" in function signature to bind "ClassType" inside a function)
example.py:28: error: Type variable "example.ClassType" is unbound
example.py:28: note: (Hint: Use "Generic[ClassType]" or "Protocol[ClassType]" base class to bind "ClassType" inside a class)
example.py:28: note: (Hint: Use "ClassType" in function signature to bind "ClassType" inside a function)
example.py:32: error: "ClassType" has no attribute "c"
Found 3 errors in 1 file (checked 1 source file)

mypy==0.942

Python 3.8.6

Am I incorrectly using TypeVar here?
I also tried using Generic as suggested in an answer here, defining the BaseClass like bellow, but it didn’t solve the issue.

ClassType = TypeVar("ClassType", bound="BaseClass")

@dataclass
class BaseClass(Generic[ClassType]):
    a: int
    b: int

I also tried without using bound attribute, as suggested in the documentation, but I get the same errors.

Is there a different way to approach this structure?

Asked By: user2586766

||

Answers:

To get rid of the mypy errors in 20 and 28, you can use generics. However, you have to make BaseExample generic and not BaseClass.

@dataclass
class BaseExample(Generic[ClassType]):
    list_attr: List[ClassType]  # line 20

    def validate(self, attrs: List[ClassType]) -> List[ClassType]:
        return [attr for attr in attrs if attr.a > 0]

However, to get rid of the error in 32, you would have to define a new type var that is bound by DerivedClass. Why? Because the attribute c is not available in BaseClass but is introduced with DerivedClass. So using BaseClass together with DerivedExample would never work.

ClassType2 = TypeVar("ClassType2", bound="DerivedClass")
@dataclass
class DerivedExample(Generic[ClassType2]):
    list_attr: List[ClassType2]  # Line 28  # Is actually List[DerivedClass]
    other_attr: int

    def select(self, attrs: List[ClassType2]) -> List[ClassType2]:
        return [attr for attr in attrs if attr.c] # line 32

You can try it out here.

Answered By: Simon Hawe