How to annotate results of iteration?

Question

from __future__ import annotations

from typing import TypeVar, Generic, Type


T = TypeVar('T')


class ListElem(Generic[T]):
    def __init__(self, value: T, nxt: ListElem[T] = None):
        self.value = value
        self.nxt = nxt


class LinkedList(Generic[T]):
    def __init__(self, elem_factory: Type[ListElem], head: ListElem[T] = None):
        self._elem_factory = elem_factory
        self._head = head

    def add(self, value: T):
        elem = self._elem_factory(value, self._head)
        self._head = elem

    def __iter__(self):
        self._next = self._head
        return self

    def __next__(self) -> ListElem[T]:
        if not self._next:
            raise StopIteration

        result = self._next
        self._next = self._next.nxt

        return result


def main():
    lst: LinkedList[int] = LinkedList(ListElem)
    lst.add(5)

    for elem in lst:
        print(elem.value)  # PyCharm see elem like int, not like ListElem


if __name__ == '__main__':
    main()

This annotating __next__ doesnt helps me

How can I annotate that LinkedList return ListElem[T] while iteration, not just T?

If I do it like this

lst: LinkedList[ListElem[int]]

I cannot annotate this int to method

def add(self, value: ?): # T = ListElem[int], but I wanna annotate value just like int
    pass

I dont wanna annotate LinkedList like

lst: LinkedList[ListElem[int], int]

Because int just repeating and annotate same - type of value inside ListElem.value

Answer 1

I would say this is a bug in PyCharm's static type checker.

That elem type should be inferred as ListElem[int] and is correctly inferred as such by mypy for example.

Your code still had a few issues. Mostly type-safety related (since you are already dealing with annotations), but a few other optimizations. One of which incidentally also fixes that PyCharm problem.

The most important one IMO is that you did not use the type parameter of ListElem in the elem_factory annotation. If you do, you'll be able to immediately bind the type argument upon initialization by passing a specified ListElem[int] class. Then you don't need to explicitly annotate lst in your main function. This also happens to satisfy/silence PyCharm.

Here is a version with my proposed changes that passes mypy --strict and should work as intended:

from __future__ import annotations
from typing import Generic, Optional, TypeVar


T = TypeVar("T")


class ListElem(Generic[T]):
    def __init__(self, value: T, nxt: Optional[ListElem[T]] = None) -> None:
        self.value = value
        self.nxt = nxt


class LinkedList(Generic[T]):
    def __init__(self, elem_factory: type[ListElem[T]], head: Optional[ListElem[T]] = None) -> None:
        self._elem_factory = elem_factory
        self._head = head

    def add(self, value: T) -> None:
        self._head = self._elem_factory(value, self._head)

    def __iter__(self) -> LinkedList[T]:
        self._next = self._head
        return self

    def __next__(self) -> ListElem[T]:
        if self._next is None:
            raise StopIteration
        result = self._next
        self._next = self._next.nxt
        return result


def main() -> None:
    lst = LinkedList(ListElem[int])
    lst.add(5)
    for elem in lst:
        print(elem.value)


if __name__ == "__main__":
    main()

However it seems that PyCharm still does not identify elem as the correct type. But at least it is not complaining.