itertools.takewhile within a generator function - why is it evaluated once only?

Question

I have a text file like this:

11
2
3
4

11

111

Using Python 2.7, I want to turn it into a list of lists of lines, where line breaks divide items in the inner list and empty lines divide items in the outer list. Like so:

[["11","2","3","4"],["11"],["111"]]

And for this purpose, I wrote a generator function that would yield the inner lists one at a time once passed an open file object:

def readParag(fileObj):
    currentParag = []
    for line in fileObj:
        stripped = line.rstrip()
    if len(stripped) > 0: currentParag.append(stripped)
    elif len(currentParag) > 0:
        yield currentParag
        currentParag = []

That works fine, and I can call it from within a list comprehension, producing the desired result. However, it subsequently occurred to me that I might be able to do the same thing more concisely using itertools.takewhile (with a view to rewriting the generator function as a generator expression, but we'll leave that for now). This is what I tried:

from itertools import takewhile    
def readParag(fileObj):
    yield [ln.rstrip() for ln in takewhile(lambda line: line != "\n", fileObj)]

In this case, the resulting generator yields only one result (the expected first one, i.e. ["11","2","3","4"]). I had hoped that calling its next method again would cause it to evaluate takewhile(lambda line: line != "\n", fileObj) again on the remainder of the file, thus leading it to yield another list. But no: I got a StopIteration instead. So I surmised that the take while expression was being evaluated once only, at the time when the generator object was created, and not each time I called the resultant generator object's next method.

This supposition made me wonder what would happen if I called the generator function again. The result was that it created a new generator object that also yielded a single result (the expected second one, i.e. ["11"]) before throwing a StopIteration back at me. So in fact, writing this as a generator function effectively gives the same result as if I'd written it as an ordinary function and returned the list instead of yielding it.

I guess I could solve this problem by creating my own class to use instead of a generator (as in John Millikin's answer to this question). But the point is that I was hoping to write something more concise than my original generator function (possibly even a generator expression). Can somebody tell me what I'm doing wrong, and how to get it right?

Answer 1

What you're trying to do is a perfect job for groupby:

from itertools import groupby

def read_parag(filename):
    with open(filename) as f:
        for k,g in groupby((line.strip() for line in f), bool):
            if k:
                yield list(g)

which will give:

>>> list(read_parag('myfile.txt')
[['11', '2', '3', '4'], ['11'], ['111']]

Or in one line:

[list(g) for k,g in groupby((line.strip() for line in open('myfile.txt')), bool) if k]

Answer 2

The other answers do a good job of explaining what is going on here, you need to call takewhile multiple times which your current generator does not do. Here is a fairly concise way to get the behavior you want using the built-in iter() function with a sentinel argument:

from itertools import takewhile

def readParag(fileObj):
    cond = lambda line: line != "\n"
    return iter(lambda: [ln.rstrip() for ln in takewhile(cond, fileObj)], [])

Answer 3

This is exactly how .takewhile() should behave. While the condition is true, it'll return elements from the underlying iterable, and as soon as it's false, it permamently switches to the iteration-done stage.

Note that this is how iterators must behave; raising StopIteration means just that, stop iterating over me, I am done.

From the python glossary on "iterator":

An object representing a stream of data. Repeated calls to the iterator’s next() method return successive items in the stream. When no more data are available a StopIteration exception is raised instead. At this point, the iterator object is exhausted and any further calls to its next() method just raise StopIteration again.

You could combine takewhile with tee to see if there are any more results in the next batch:

import itertools

def readParag(filename):
    with open(filename) as f:
        while True:
            paras = itertools.takewhile(lambda l: l.strip(), f)
            test, paras = itertools.tee(paras)
            test.next()  # raises StopIteration when the file is done
            yield (l.strip() for l in paras)

This yields generators, so each item yielded is itself a generator. You do need to consume all elements in these generators for this to continue to work; the same is true for the groupby method listed in another answer.

Answer 4

If the file contents fit into memory, there is a much easier way to get the groups separated by blank lines:

with open("filename") as f:
    groups = [group.split() for group in f.read().split("\n\n")]

This approach can be made more robust by using re.split() instead of str.split() and by filtering out potential empty groups resulting from four or more consecutive line breaks.

Answer 5

This is the documented behavior of takewhile. It takes while the condition is true. It doesn't start up again if the condition later becomes true again.

The simple fix is to make your function just call takewhile in a loop, stopping when takewhile has nothing more to return (i.e., at the end of the file):

def readParag(fileObj):
    while True:      
        nextList = [ln.rstrip() for ln in takewhile(lambda line: line != "\n", fileObj)]
        if not nextList:
            break
        yield nextList

Answer 6

You can call takewhile multiple times:

>>> def readParagGenerator(fileObj):
...     group = [ln.rstrip() for ln in takewhile(lambda line: line != "\n", fileObj)]
...     while len(group) > 0:
...         yield group
...         group = [ln.rstrip() for ln in takewhile(lambda line: line != "\n", fileObj)]
... 
>>> list(readParagGenerator(StringIO(F)))
[['11', '2', '3', '4'], ['11'], ['111']]