Iteration: iter, next and lazy generators

A for loop works on lists, strings, dicts, files and much more, and it does it with just two built-in functions. This loop:

for colour in ["red", "amber", "green"]:
    print(colour)

does the same as this:

lights = iter(["red", "amber", "green"])
while True:
    try:
        colour = next(lights)
    except StopIteration:
        break
    print(colour)

• iter(x) asks x for an iterator, an object that hands out items one at a time;
• next(it) asks the iterator for the next item;
• when there are none left, next raises StopIteration, and the loop ends quietly.

An iterable is anything iter() works on: a list is iterable, but it isn't an iterator itself. An iterator remembers where it is, and it gets used up:

letters = iter("abc")
print(next(letters))   # a
print(list(letters))   # ['b', 'c']: the rest
print(list(letters))   # []: nothing left

Writing an iterator

A class becomes an iterator with two special methods: __next__ gives the next item (or raises StopIteration), and __iter__ returns the iterator itself, so it works in a for loop too:

class Countdown:
    def __init__(self, start):
        self.current = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.current <= 0:
            raise StopIteration
        self.current -= 1
        return self.current + 1


print(list(Countdown(3)))   # [3, 2, 1]

Generators: the short way

A function with yield in it is a generator function. Calling it doesn't run its body: it returns a generator, an iterator that runs the body a step at a time. Each next() runs until the next yield, hands that value over and pauses right there, with its local variables kept. The prints show you exactly when each part runs:

def slow_squares(n):
    print("  (starting)")
    for i in range(n):
        print(f"  (working out {i} squared)")
        yield i * i
    print("  (finished)")


gen = slow_squares(3)
print("Made the generator. Nothing has run yet.")
print("first:", next(gen))
for value in gen:
    print("loop got", value)

Made the generator. Nothing has run yet.
  (starting)
  (working out 0 squared)
first: 0
  (working out 1 squared)
loop got 1
  (working out 2 squared)
loop got 4
  (finished)

That's what lazy means: no value is worked out before it's asked for. A generator can describe a million items, or an endless stream, while holding only one at a time. A generator expression is the one-line form, written like a list comprehension with round brackets:

total = sum(n * n for n in range(1_000_000))   # no million-item list is built
print(total)

More in the Python docs: iterator types, generators and the yield expression.