Python Iterators and Generators

When you write a for loop over a list, Python takes the items one by one. This one-by-one working is done by an iterator. A generator is an easy way to create your own iterator using the yield keyword.


Iterable and iterator

  • An iterable is anything you can loop over, like a list, tuple, string, dictionary or set.
  • An iterator is an object that gives the next item each time you ask. You get it from an iterable with iter(), and ask for the next item with next().
fruits = ["apple", "banana", "cherry"]
it = iter(fruits)

print(next(it))
print(next(it))
print(next(it))
Output of the above program is

apple
banana
cherry

If you call next(it) a fourth time, there are no items left and Python raises StopIteration. A for loop catches this error silently and stops. You can also read the interview question What are iterators in Python?


Create your own iterator

A class becomes an iterator when it has two methods: __iter__() that returns the iterator itself, and __next__() that returns the next value or raises StopIteration.

class CountUpTo:
    def __init__(self, limit):
        self.limit = limit
        self.current = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.limit:
            raise StopIteration
        self.current += 1
        return self.current

for number in CountUpTo(3):
    print(number)
Output of the above program is

1
2
3


Generators

Writing a class for every iterator is long. A generator function is a normal function that uses yield instead of return. Each time it reaches yield, it gives a value and pauses. The next call continues from the same place.

def count_up(limit):
    current = 1
    while current <= limit:
        yield current
        current += 1

for number in count_up(3):
    print(number)
Output of the above program is

1
2
3

You can also take the values by hand with next().

gen = count_up(2)
print(next(gen))
print(next(gen))
Output of the above program is

1
2

Fibonacci generator

def fibonacci(limit):
    a, b = 0, 1
    while a < limit:
        yield a
        a, b = b, a + b

print(list(fibonacci(30)))
Output of the above program is

[0, 1, 1, 2, 3, 5, 8, 13, 21]


Generator expressions

A generator expression looks like a list comprehension but uses round brackets. It creates the values one at a time, and does not build the full list in memory.

squares = (x * x for x in range(5))

print(next(squares))
print(next(squares))
print(sum(squares))
Output of the above program is

0
1
29

The first two values (0 and 1) were taken by next(). sum() then added only the remaining ones: 4 + 9 + 16 = 29. A generator can be used only once, so create it again if you need the values again.


Iterator vs generator

PointIterator classGenerator function
How to writeClass with __iter__() and __next__()Function with yield
Code sizeLongerShort
StateYou save it in attributesPython remembers it for you
Use whenYou need a full object with extra methodsYou just need to produce a sequence of values
Note: Generators are great for big files or endless sequences, because only one value stays in memory at a time. Read the interview question What is a generator in Python? for more.

Summary

  • An iterable can be looped over. iter() gives its iterator and next() gives the next item.
  • A custom iterator class needs __iter__() and __next__().
  • A function with yield is a generator. It pauses and continues later.
  • Generator expressions use round brackets.
  • Generators save memory, but you can use them only once.

Congratulations! Chapter Finished. Learn more about the similar topics:
Exercises & Assignments
No Content Found.
Interview Questions & Answers
No Content Found.