Tag: Generators

Python – Generators

Do you know the difference between the following lines in Python?

square_numbers_1 = [x*x for x in [1,2,3,4,5,6]]
square_numbers_2 = (x*x for x in [1,2,3,4,5,6])

The first line will create a list of squares for each element of [1,2,3,4,5,6] in memory which means you can access it by using square brackets:

# This will return 1
print square_numbers_1[0]

The second line will return a Generator object which means you cannot use the square brackets.
You can access the values using a for loop or the next() method.

# This will print all items:
# 1
# 4
# 9
# 16
# 25
# 36
for number in square_numbers_2:
    print number 

# This will return one item each time.
# The following line will return 1.
print next(square_numbers_2)

You need to call ‘next’ whenever you need an item and you will receive an exception (StopIteration) when you call it more than the number of items. More than six times for square_numbers_2.

You can also create a Generator by calling yield as in the following example:

def square_numbers(numbers):
    for n in numbers:
        yield (n*n)

square_numbers_generator = square_numbers([1,2,3,4,5,6])

for number in square_numbers_generator:
    print number

Why is Generators useful?

  • Reduce the amount of memory used as it will only allocate memory when you retrieve an item
  • Run faster because it stops when it gets to the first yield in the above example