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

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?

CodePudding user response：

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.