Python subclass of abstract baseclass shows type error

Question

pylance gives an type error while the types of the variables are defined. Below a small example which shows the structure of the code and gives also the error:

from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generator, Any


@dataclass
class Base(ABC):
    data: list[str | int]

    def __iter__(self) -> Generator[str | int, Any, None]:
        for xin self.data:
            yield x

    def __getitem__(self, key) -> str | int:
        return self.data[key]

    def __len__(self) -> int:
        return len(self.data)

    @abstractmethod
    def some_general_function(self) -> None:
        pass


class Test1(Base):
    data: list[str]

    def some_general_function(self) -> None:
        print("General function from list with strings")


if __name__ == '__main__':
    test = Test1(['a', 'b', 'abc'])
    for ele in test:
        print(ele.capitalize)

The error is on the last line, the error says that the function .capitalize is not available for the type int.

In the implementation of the class (Test1) it is defined that the list data consist of elements with type str.

Also when this is tested in a __post_init__ method in the Test1 class the error still exists.

class Test1(Base):
    data: list[str]

    def __post_init__(self):
        if not any(map(lambda x: isinstance(x, str), self.data)):
            raise TypeError("data should be a `str`")

The only solution is to rewrite the __iter__ of the baseclass, since the type hint there says that the generator can be a str and an int.

Is it possible to only rewrite the the return type (or yield type) of a function without rewriting the function itself?

Answer 1

With the TypeVar and Generic from the library typing it is possible to adjust the code so it returns what one wants.

from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generator, Any, TypeVar, Generic


_T = TypeVar("_T")

@dataclass
class Base(ABC, Generic[_T]):
    data: list[_T]

    def __iter__(self) -> Generator[_T, Any, None]:
        for x in self.data:
            yield x

    def __getitem__(self, key: int) -> _T:
        return self.data[key]

    def __len__(self) -> int:
        return len(self.data)

    @abstractmethod
    def some_general_function(self) -> None:
        pass


class Test1(Base[str]):
    def some_general_function(self) -> None:
        print("General function from list with strings")


if __name__ == '__main__':
    test = Test1(['a', 'b', 'abc'])
    for ele in test:
        print(ele.capitalize())

An explanation about TypeVar can be found here: https://dev.to/decorator_factory/typevars-explained-hmo

And the Generic parameter is explained here: https://mypy.readthedocs.io/en/stable/generics.html