Python · iterators & generators

Generators

Functions that pause: yield hands out one value at a time and picks up where it left off.

A function that can pause

Calling it runs nothing

Calling a generator function only creates a generator object, with its frame parked and its body not yet started.

next() resumes

The frame goes back on the stack and runs until the next yield.

yield pauses

The value goes to the caller and the frame is parked again, local variables and all. Watch the Paused generators shelf.

Generator function basics


          # generator with yield
          def count_up_to(n):
              count = 1
              while count <= n:
                  yield count
                  count += 1

          for num in count_up_to(3):
              print(num)
          # ans: 1, 2, 3

Yield vs return


          # yield pauses function
          def simple_gen():
              yield 1
              yield 2
              yield 3

          gen = simple_gen()
          next(gen)
          # ans: 1
          next(gen)
          # ans: 2

return would end the call for good; yield keeps the frame for the next next().

A generator is an iterator


          # generator is an iterator
          def my_gen():
              yield 1
              yield 2

          gen = my_gen()
          next(gen)
          # ans: 1
          # can only iterate once
          list(gen)
          # ans: [2]
          list(gen)
          # ans: []

Infinite generator


          # infinite sequence
          def infinite_counter():
              count = 0
              while True:
                  yield count
                  count += 1

          # use with break or take first n
          counter = infinite_counter()
          for i, num in enumerate(counter):
              if i >= 3:
                  break
              print(num)
          # ans: 0, 1, 2

An endless loop is fine in a generator: it only runs when asked for the next value.

Generator with a condition


        # filter in generator
        def even_numbers(n):
            for i in range(n):
                if i % 2 == 0:
                    yield i

        list(even_numbers(10))
        # ans: [0, 2, 4, 6, 8]

Only values that pass the test are yielded; the others are skipped inside the loop.

yield from


          # delegate to another generator
          def generator1():
              yield 1
              yield 2

          def generator2():
              yield from generator1()
              yield 3

          list(generator2())
          # ans: [1, 2, 3]

yield from passes on every value of the inner generator, then carries on.

Exercises

Predict the answer, then press → to check.

Exercise 1


          # what is returned?
          def gen():
              yield 1

          result = gen()
          type(result)
          # ans: <class 'generator'>

Exercise 2


        # yield in loop?
        def gen():
            for i in range(3):
                yield i

        list(gen())
        # ans: [0, 1, 2]

Exercise 3


          # generator exhaustion?
          def gen():
              yield 1
              yield 2

          g = gen()
          next(g)
          # ans: 1
          next(g)
          # ans: 2
          next(g)
          # ans: StopIteration

Exercise 4


        # multiple yields?
        def gen():
            yield 1
            yield 2
            return 3

        list(gen())
        # ans: [1, 2] (the return value is ignored)

Exercise 5


        # generator vs list?
        def gen():
            yield 1

        list(gen())
        # ans: [1]
        gen()
        # ans: <generator object>

Exercise 6


        # fibonacci generator?
        def fib(n):
            a, b = 0, 1
            for _ in range(n):
                yield a
                a, b = b, a + b

        list(fib(5))
        # ans: [0, 1, 1, 2, 3]

Exercise 7


        # yield from list?
        def gen():
            yield from [1, 2, 3]

        list(gen())
        # ans: [1, 2, 3]

Exercise 8


          # generator with return?
          def gen():
              yield 1
              return 2

          g = gen()
          next(g)
          # ans: 1
          next(g)
          # ans: StopIteration (with value 2)

Exercise 9


          # send to generator?
          def gen():
              value = yield
              print(value)

          g = gen()
          next(g)  # prime
          g.send("Hello")
          # ans: prints "Hello"

send() resumes the generator and makes the paused yield evaluate to the value sent.

Exercise 10


        # generator memory?
        def gen():
            for i in range(1000000):
                yield i
        # ans: memory efficient (lazy)

        lst = [i for i in range(1000000)]
        # ans: all in memory at once

Remember

Lazy

Calling a generator function runs nothing until the first next().

Pause and resume

Each yield hands out a value and keeps the local variables for later.

One pass only

Once exhausted, a generator raises StopIteration and gives nothing more.

Small memory

Values are made one at a time instead of stored all at once.