Question

I’m writing a Python generator which looks like “cat”. My specific use case is for a “grep like” operation. I want it to be able to break out of the generator if a condition is met:

summary={}
for fn in cat("filelist.dat"):
    for line in cat(fn):
        if line.startswith("FOO"):
            summary[fn] = line
            break

So when break happens, I need the cat() generator to finish and close the file handle to fn.

I have to read 100k files with 30 GB of total data, and the FOO keyword happens in the header region, so it is important in this case that the cat() function stops reading the file ASAP.

There are other ways I can solve this problem, but I’m still interested to know how to get an early exit from a generator which has open file handles. Perhaps Python cleans them up right away and closes them when the generator is garbage collected?

Thanks,

Ian

Answer 1

By implementing the context protocol and the iterator protocol in the same object, you can write pretty sweet code like this:

with cat("/etc/passwd") as lines:
    for line in lines:
        if "mail" in line:
            print line.strip()
            break

This is a sample implementation, tested with Python 2.5 on a Linux box. It reads the lines of /etc/passwd until it finds the one for user audio, and then stops:

from __future__ import with_statement


class cat(object):

    def __init__(self, fname):
        self.fname = fname

    def __enter__(self):
        print "[Opening file %s]" % (self.fname,)
        self.file_obj = open(self.fname, "rt")
        return self

    def __exit__(self, *exc_info):
        print "[Closing file %s]" % (self.fname,)
        self.file_obj.close()

    def __iter__(self):
        return self

    def next(self):
        line = self.file_obj.next().strip()
        print "[Read: %s]" % (line,)
        return line


def main():
    with cat("/etc/passwd") as lines:
        for line in lines:
            if "mail" in line:
                print line.strip()
                break


if __name__ == "__main__":
    import sys
    sys.exit(main())

Or even simpler:

with open("/etc/passwd", "rt") as f:
    for line in f:
        if "mail" in line:
            break

File objects implement the iterator protocol (see http://docs.python.org/library/stdtypes.html#file-objects)