Python Cookbook

Wrapping an Unbounded Iterator to Restrict Its Output

Credit: Tom Good

Problem

You need to filter the sequence produced by a potentially unbounded Python 2.2 iterator or limit the sequence length by a condition.

Solution

Python 2.2 generators are suitable for wrapping other generators (or other kinds of iterators) and tweaking their output—for example, by limiting the output’s length:

from _ _future_ _ import generators

def genWhile(g, condition):
    """ Run generator g, stopping when condition(g.next(  )) is false. condition
    can be any callable. genWhile returns an iterator. """
    g = iter(g)
    while 1:
        next = g.next(  )
        if condition(next):
            yield next
        else:
            return

def take(n, g):
    """ A subiterator limited to the first n items of g's sequence """
    g = iter(g)
    for i in range(n): yield g.next(  )

def drop(n, g):
    """ A subiterator removing the first n items from g's sequence """
    g = iter(g)
    for i in range(n): g.next(  )
    while 1: yield g.next(  )

# an example of an unbounded sequence generator
def genEven(  ):
    x = 0
    while 1:
        x += 2
        yield x

def main(  ):
    print [x for x in genWhile(genEven(  ), lambda x: x<12)]
    print [x for x in take(5, genEven(  ))]
    print [x for x in take(5, drop(5, genEven(  )))]

Discussion

With Python 2.2 and later, you can make iterators that return unbounded output (for example, see Recipe 17.11). By creating a wrapper generator that runs another iterator, you can restrict the resulting sequence to a defined subset. The g=iter(g) idiom at the start of each wrapper in this recipe ensures ...