Iterators & Generators -2

Generators

Generators simplifies creation of iterators. A generator is a function that produces a sequence of results instead of a single value.

def yrange(n):
    i = 0
    while i < n:
        yield i
        i += 1

Each time the yield statement is executed the function generates a new value.

>>> y = yrange(3)
>>> y
<generator object yrange at 0x401f30>
>>> y.next()
0
>>> y.next()
1
>>> y.next()
2
>>> y.next()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

So a generator is also an iterator. You don’t have to worry about the iterator protocol.
The word “generator” is confusingly used to mean both the function that generates and what it generates. In this chapter, I’ll use the word “generator” to mean the genearted object and “generator function” to mean the function that generates it.
Can you think about how it is working internally?
When a generator function is called, it returns an generator object without even beginning execution of the function. When next` method is called for the first time, the function starts executing until it reaches yield statement. The yielded value is returned by the next call.
The following example demonstrates the interplay between yield and call to next method on generator object.

>>> def foo():
...     print "begin"
...     for i in range(3):
...         print "before yield", i
...         yield i
...         print "after yield", i
...     print "end"
...
>>> f = foo()
>>> f.next()
begin
before yield 0
0
>>> f.next()
after yield 0
before yield 1
1
>>> f.next()
after yield 1
before yield 2
2
>>> f.next()
after yield 2
end
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

Lets see an example:

def integers():
    """Infinite sequence of integers."""
    i = 1
    while True:
        yield i
        i = i + 1

def squares():
    for i in integers():
        yield i * i

def take(n, seq):
    """Returns first n values from the given sequence."""
    seq = iter(seq)
    result = []
    try:
        for i in range(n):
            result.append(seq.next())
    except StopIteration:
        pass
    return result

print take(5, squares()) # prints [1, 4, 9, 16, 25]

 Generator Expressions

Generator Expressions are generator version of list comprehensions. They look like list comprehensions, but returns a generator back instead of a list.

>>> a = (x*x for x in range(10))
>>> a
<generator object <genexpr> at 0x401f08>
>>> sum(a)
285

We can use the generator expressions as arguments to various functions that consume iterators.

>>> sum(((x*x for x in range(10)))
285

When there is only one argument to the calling function, the parenthesis around generator expression can be omitted.

>>> sum(x*x for x in range(10))
285

Another fun example:
Lets say we want to find first 10 (or any n) pythogorian triplets. A triplet (x, y, z) is called pythogorian triplet if x*x + y*y == z*z.
It is easy to solve this problem if we know till what value of z to test for. But we want to find first n pythogorian triplets.

>>> pyt = ((x, y, z) for z in integers() for y in xrange(1, z) for x in range(1, y) if x*x + y*y == z*z)
>>> take(10, pyt)
[(3, 4, 5), (6, 8, 10), (5, 12, 13), (9, 12, 15), (8, 15, 17), (12, 16, 20), (15, 20, 25), (7, 24, 25), (10, 24, 26), (20, 21, 29)]

5.3.1. Example: Reading multiple files

Lets say we want to write a program that takes a list of filenames as arguments and prints contents of all those files, like cat command in unix.
The traditional way to implement it is:

def cat(filenames):
    for f in filenames:
        for line in open(f):
            print line,

Now, lets say we want to print only the line which has a particular substring, like grep command in unix.

def grep(pattern, filenames):
    for f in filenames:
        for line in open(f):
            if pattern in line:
                print line,

Both these programs have lot of code in common. It is hard to move the common part to a function. But with generators makes it possible to do it.

def readfiles(filenames):
    for f in filenames:
        for line in open(f):
            yield line

def grep(pattern, lines):
    return (line for line in lines if pattern in lines)

def printlines(lines):
    for line in lines:
        print line,

def main(pattern, filenames):
    lines = readfiles(filenames)
    lines = grep(pattern, lines)
    printlines(lines)

The code is much simpler now with each function doing one small thing. We can move all these functions into a separate module and reuse it in other programs.
Problem 2: Write a program that takes one or more filenames as arguments and prints all the lines which are longer than 40 characters.
Problem 3: Write a function findfiles that recursively descends the directory tree for the specified directory and generates paths of all the files in the tree.
Problem 4: Write a function to compute the number of python files (.py extension) in a specified directory recursively.
Problem 5: Write a function to compute the total number of lines of code in all python files in the specified directory recursively.
Problem 6: Write a function to compute the total number of lines of code, ignoring empty and comment lines, in all python files in the specified directory recursively.
Problem 7: Write a program split.py, that takes an integer n and a filename as command line arguments and splits the file into multiple small files with each having n lines.