Lazy list builders: Generators in Python

Now let's turn it into a version that uses generators. The trick is simply to replace every point at which you would add a value to the end results list and replace it with a yield statement:

def xcombination(seq,length):
if not length: yield []
else:
for i in xrange(len(seq)):
for result in xcombination(seq[i+1:], length-1):
yield [seq[i]]+result


Now we have a version that works exactly the same, except that it does the work only as you need it:

>>> comb = xcombination("ABCDE", 3)
>>> comb.next()
['A', 'B', 'C']
>>> comb.next()
['A', 'B', 'D']
>>> list(comb)
[['A', 'B', 'E'], ['A', 'C', 'D'], ['A', 'C', 'E'], ['A', 'D', 'E'], ['B', 'C', 'D'], ['B', 'C', 'E'], ['B', 'D', 'E'], ['C', 'D', 'E']]
>>> comb2 = xcombination("ABCDE",2)
>>> for i in xrange(3):
... print comb2.next()
... ['A', 'B']
['A', 'C']
['A', 'D']


In the last command, even though there are 10 different combinations of letters, only three are generated, saving us 70 percent of the computation time of the standard function.

The really useful thing about generators is that they can, for most needs, be used as a drop in replacement for list comprehensions. All you need to do is replace the brackets that would go around list comprehensions with parentheses. Take the last example in the article on list comprehensions. Rather than:

>>> guests = ['Chris', 'Brendan', 'Jimmy', 'Mel', 'Mike', 'Jess']
>>> [(seat1, seat2) for seat1 in guests for seat2 in guests if seat1 != seat2]
...


You can write:

>>> guests = ['Chris', 'Brendan', 'Jimmy', 'Mel', 'Mike', 'Jess']
>>> ((seat1, seat2) for seat1 in guests for seat2 in guests if seat1 != seat2)



And then you can use it like you would any other generator:

>>> seating = ((seat1, seat2) for seat1 in guests for seat2 in guests if seat1 != seat2)
>>> for i in xrange(10):
... print seating.next()
...
('Chris', 'Brendan')
('Chris', 'Jimmy')
('Chris', 'Mel')
('Chris', 'Mike')
('Chris', 'Jess')
('Brendan', 'Chris')
('Brendan', 'Jimmy')
('Brendan', 'Mel')
('Brendan', 'Mike')
('Brendan', 'Jess')


Up til now we've been using generators as a way of building lists to save computation time over other methods. That's great, but where they really shine is when computing the whole list would be impossible. Take the Fibonacci sequence, in which each number is the sum of the previous two. Say we want a function that generates the sequence up to a given number:

>>> def fib(n):
...     a, b = 0, 1
...     l = [a]
...     while b < n:
...             l += [b]
...             a, b = b, a+b
...     return l
... 
>>> fib(20)
[0, 1, 1, 2, 3, 5, 8, 13]

Now this works fine, so long as we want to stop counting up when we reach a certain number. If we want to change the criteria for stopping, however, we have to rewrite the function. Instead we can implement it as a generator (implementation borrowed from the python PEP on generators):

>>> def xfib():
...     a,b = 0,1
...     while True:
...             yield b
...             a, b = b, a+b
... 
>>> fibseries = xfib()
>>> b = fibseries.next()
>>> while b < 20:
...     print b
...     b = fibseries.next()
... 
1
1
2
3
5
8
13

Or, if we want to stop at the first palindrome (greater than one digit) we just need to change the loop condition:

>>> fibseries = xfib()
>>> b = fibseries.next()
>>> while b < 10 or not list(str(b)) == list(reversed(str(b))):
... print b
... b = fibseries.next()
...
1
1
2
3
5
8
13
21
34


And that's it (The next value in the series is 55). By using a generator we can keep the list building implementation separate from the logic relating to when to stop generating it, while still only calculating as many values as we need.

When should you use generators over list comprehensions? Firstly, if you're going to use the whole list anyway, it's better to use a list comprehension -- they're faster since they don't have the added overhead of a generator function call. If you're going to use the first part of the list, use a generator, since you'll save CPU time.