CODEWITHSUNDEEP
pythonpython-typingmypy
David Michael Gang
David Michael Gang

Reputation: 7309

Confusion of annotating generator function as iterator

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:

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?

Upvotes: 12

Views: 15031

Answers (1)

ethanhs
ethanhs

Reputation: 1843

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.

Upvotes: 18

Related Questions