I have a BaseList
container that takes BaseItem
‘s as items. I then derive a new list CustomList
and I want it to hold CustomItem
‘s.
How to I type methods in BaseList
many methods to accept BaseItem
. And also tell it to use CustomItem
the derived list CustomList
?
Here is the code I have so far:
from __future__ import annotations
from typing import List, Union, Iterator, overload
class BaseItem:
pass
class BaseList:
def __init__(self) -> None:
self.items: List[BaseItem] = []
def add_item(self, item: BaseItem) -> None:
self.items.append(item)
def __iter__(self) -> Iterator[BaseItem]:
return iter(self.items)
class CustomItem(BaseItem):
def __init__(self, value: int) -> None:
self.value: int = value
class CustomList(BaseList):
def sum(self) -> int:
# mypy error: "BaseItem" has no attribute "value"
return sum(item.value for item in self)
# no error, but it should say: incompatible type; expected "CustomItem".
CustomList().add_item(BaseItem())
container = CustomList()
container.add_item(CustomItem(10))
for item in container:
# mypy error: `"BaseItem" has no attribute "value"`
print(item.value)
Question
How can I define a type A
and say in BaseList
it should resolve to BaseItem
and in CustomList
it should resolve to CustomItem
?
Tried so far
I tried using TypeVar with the bounds attribute, but I didn’t seem to get it work.
Constraints
I have more than just one subclass so I don’t want to re-implement all these methods with different types. Also there are many more of these methods with similar problems, like __contains__
, __getitem__
, class specific items, etc.
The type hints should be also resolved by PyCharm, as the goal is to get correct code-complete there.
Advertisement
Answer
Make BaseList class Generic:
class BaseItem:
pass
T = TypeVar('T', bound=BaseItem)
class BaseList(Generic[T]):
def __init__(self) -> None:
self.items: List[T] = []
def add_item(self, item: T) -> None:
self.items.append(item)
def __iter__(self) -> Iterator[T]:
return iter(self.items)
class CustomItem(BaseItem):
def __init__(self, value: int) -> None:
self.value: int = value
class CustomList(BaseList[CustomItem]):
def sum(self) -> int:
reveal_type(self.items) # Mypy: Revealed type is "builtins.list[CustomItem*]"
return sum(item.value for item in self)