Search code examples
pythontypingpylance

Python subclass of abstract baseclass shows type error


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?


Solution

  • 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