How to type annotate a chunker function?

Question

I try the following:

from typing import Iterator, List, Any, TypeVar
T = TypeVar('T')
def chunker(seq: Iterator[T], size: int) -> Iterator[List[T]]:
    for i in seq:
        yield []
for i in chunker([1, 2, 3, 4], 5):
    pass

However, pyright tells me that:

  test.py:5:18 - error: Argument of type "list[int]" cannot be assigned to parameter "seq" of type "Iterator[T@chunker]" in function "chunker"
    "list[int]" is incompatible with protocol "Iterator[T@chunker]"
      "__next__" is not present (reportGeneralTypeIssues)
1 error, 0 warnings, 0 informations 

How should I annotate a function that takes something that you can for i in seq it?

I tried searching on stack and google, but I guess like my google-fu is not good enough. I could use Union[Iterator[T], List[T]] but I think there should be something better?

Answer 1

An iterator is an object with state that remembers where it is during iteration. It has a __next__ method that returns the next value in the iteration, updates the state to point at the next value and signals when it is done by raising StopIteration.

As you can see, a list does not have a __next__ method:

>>> [1,2,3].__next__()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'list' object has no attribute '__next__'

Instead, what you really want here is an iterable. An iterable is:

• anything that can be looped over (i.e. you can loop over a string or file)
• anything that can appear on the right-side of a for-loop: for x in iterable: ...
• anything you can call with iter() that will return an iterator: iter(obj)
• an object that defines __iter__ that returns a fresh iterator

A list is indeed an iterable:

>>> [1,2,3].__iter__()
<list_iterator object at 0x103e68310>

More info can be found in this question: What are iterator, iterable, and iteration?

So you should rewrite your example as:

from collections.abc import Generator, Iterable
from typing import Any, TypeVar

T = TypeVar('T')
def chunker(seq: Iterable[T], size: int) -> Generator[list[T], None, None]:
    for i in seq:
        yield []
for i in chunker([1, 2, 3, 4], 5):
    pass

Notice that since Python version 3.9, typing.Iterable is deprecated. More details can be found here. Also, typing.List is a deprecated alias of list (more details here).

A simpler way is to use Iterator instead of Generator (every generator is an iterator!!):

from collections.abc import Iterable, Iterator
from typing import Any, TypeVar

T = TypeVar('T')
def chunker(seq: Iterable[T], size: int) -> Iterator[list[T]]:
    for i in seq:
        yield []
for i in chunker([1, 2, 3, 4], 5):
    pass