confusion of annotating generator function as iterator

Question

In the python typing documentation it is written:

Alternatively, annotate your generator as having a return type of either Iterable[YieldType] or Iterator[YieldType]:

def infinite_stream(start: int) -> Iterator[int]:
    while True:
        yield start
        start += 1

I wrote a very easy example of printing an infinite stream. I have a generator function which is passed to another function and then called.

from typing import Iterator


def infinite_stream(start: int) -> Iterator[int]:
    while True:
        yield start
        start += 1


def print_infinite_stream(inf_iterator: Iterator[int]):
    for x in inf_iterator(5):
        print(x)


print_infinite_stream(infinite_stream)

With mypy i get two errors:

• error:Iterator[int] not callable

• error:Argument 1 to "print_infinite_stream" has incompatible type "Callable[[int], Iterator[int]]"; expected "Iterator[int]"

I am confused why i am getting these errors as i worked according to the documentation and have the latest python (3.6.5) and mypy (0.590) installed. What is wrong here?

Answer 1

annotate your generator as having a return type of either Iterable[YieldType] or Iterator[YieldType]

Generator functions return generators, they are not generators themselves. If you do:

reveal_type(infinite_stream), you will get something like Callable[[int], Iterator[int]].

What you want is the return value of the function, the actual iterator.

from typing import Iterator


def infinite_stream(start: int) -> Iterator[int]:
    while True:
        yield start
        start += 1


def print_infinite_stream(inf_iterator: Iterator[int]):
    for x in inf_iterator:
        print(x)


print_infinite_stream(infinite_stream(5))

This makes more sense as now print_infinite_stream handles any iterator, not just your generator function. If you reveal_type(infinite_stream(5)) you should get something like Iterator[int], which is exactly what you want.