Generators

We will look at the generators in python using a running example which is to merge two sorted, sentinel-terminated lists. (Using sentinel-terminated lists simplifies the merge code; using null-terminated lists is more common but requires more code. The difference between the two is not important for today's discussion.

The first version of the code is C-ish: it uses subscripting. (As someone pointed out in class, a real C programmer would probably use pointers and the code would be somewhat less ugly.)

inf = 1000000
x = [2, 4, 6, 10, 20, 24, 36, inf]
y = [1, 3, 5, 15, 27, 45, 57, inf]

def mergelists1(l1,l2):
    out = []
    i1 = 0; i2 = 0
    while l1[i1]<>inf or l2[i2]<>inf:
        if l1[i1]<=l2[i2]:
            out.append(l1[i1])
            i1 = i1+1
        else:
            out.append(l2[i2])
            i2 = i2+1
    # we want to process both list fully. This means that
    # at the end we want
    # l1[i1]==inf and l2[i2]==inf
    # Remember that the proof rule for loops tells us therefore
    # that the loop test needs to be 
    # l1[i1]<>inf or l2[i2]<>inf
    return out

print mergelists1(x,y)

The code of mergelists2 takes advantage of iterators and is an elegant expression of the merge algorithm.

def mergelists2(l1,l2):
    out = []
    it1 = iter(l1); it2 = iter(l2)
    e1 = it1.next(); e2 = it2.next()
    while e1<>inf or e2<>inf:
        if e1<e2:
            out.append(e1)
            e1 = it1.next()
        else:
            out.append(e2)
            e2 = it2.next()
    return out

You might now ask, "what if I want to build an iterator for my own data structures?" Let's examine the case of lists first and try to reproduce the behavior of iter(). An iterator is an object with a next method. In order to produce the elements of the list one by one we need to maintain some state in the object saying what was the last element produced.

class iter:
    def __init__(self,l):
        # we choose to represent "none produced yet" with index -1
        self.pos = -1
        self.l =l
    def next(self):
        # choosing -1 for "none produced yet" makes next() easy
        self.pos = self.pos+1
        return self.l[self.pos]

Notice how in class iter we had to choose a representation for the state of the iteration and manage it by hand. Compare how it is done in iter with how we would naturally write a loop to examine the elements of a list one by one:

   for e in l:
      ... do something with e

A python generator is any function that contains a yield statement. When such a function is called a new object containing a next() method is created, all the parameters are remembered and the object is returned to the function's caller. When the next() method is called the body of the defined function begins to execute and continues up the first yield statement executed. yield is like return in that the caller sees the value yielded as the value of its call to next(). However, unlike return the state of the function is remembered and execution will resume at that point when next() is called again. Eventually there may be no more values to produce as the result of next at which point the function can return at which point the python runtime raises the StopIteration exception.

We can now return to our list iterator implementation and see how it would be done with a generator.

def iter(l):
    for e in l:
        yield e

The hand-crafted iterator class for lists wasn't too painful but consider now what you would have to do to build an iterator for a binary tree class such as

# A binary tree class.
class Tree:

    def __init__(self, label, left=None, right=None):
        self.label = label
        self.left = left
        self.right = right
        
    # __repr__ gives a string representation of any object
    # used, e.g.,  to print the object
    def __repr__(self, level=0, indent="    "):
        s = level*indent + `self.label`
        if self.left:
            s = s + "\n" + self.left.__repr__(level+1, indent)
        if self.right:
            s = s + "\n" + self.right.__repr__(level+1, indent)
        return s

# Create a Tree from a list.
def tree(list):
    n = len(list)
    if n == 0:
        return None
    i = n / 2
    return Tree(list[i], tree(list[:i]), tree(list[i+1:]))

    
    def iter(t):
        if t:
            if t.left:
                for x in t.left.iter():
                    yield x
            yield t.label
            # t.right() also produces an iterator because
            # of the __iter__ method that is automagically called
            # when an iterator is needed for a class
            if t.right:
                for x in t.right.iter():
                    yield x

def mergetrees(t1,t2):
    out = []
    it1 = t1.iter()
    it2 = t2.iter()
    e1 = it1.next(); e2 = it2.next()
    while e1<>inf or e2<>inf:
        if e1<e2:
            out.append(e1)
            e1 = it1.next()
        else:
            out.append(e2)
            e2 = it2.next()
    return out

xt = tree(x)
yt = tree(y)
print mergetrees(xt, yt)