如何从带有打开文件句柄的Python生成器中退出

Question

我正在写一个Python生成器，功能类似于“cat”命令。我的具体需求是进行类似“grep”的操作。我希望在满足某个条件时，能够退出这个生成器：

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

所以当触发break时，我需要cat()这个生成器结束，并且关闭对fn的文件句柄。

我需要读取10万个文件，总共30GB的数据，而FOO这个关键词出现在文件的开头部分，所以在这种情况下，cat()函数需要尽快停止读取文件。

虽然我有其他方法可以解决这个问题，但我还是想知道如何从一个打开文件句柄的生成器中提前退出。也许Python会在生成器被垃圾回收时立即清理并关闭这些文件句柄？

谢谢，

伊恩

Answer 1

请考虑这个例子：

def itertest():
    try:
        for i in xrange(1000):
            print i
            yield i
    finally:
        print 'finally'

x = itertest()

for i in x:
    if i > 2:
        break

print 'del x'
del x

print 'exit'

0
1
2
3
del x
finally
exit

这个例子展示了在迭代器被清理后，finally 语句会被执行。我认为 __del__(self) 是在调用 self.close()，你可以在这里查看更多信息：https://docs.python.org/2.7/reference/expressions.html#generator.close

Answer 2

生成器有一个叫做 close 的方法，当你在 yield 语句处调用它时，会抛出一个 GeneratorExit 的异常。如果你专门捕捉这个异常，就可以运行一些清理代码：

import contextlib
with contextlib.closing( cat( fn ) ):
    ...

然后在 cat 函数里：

try:
    ...
except GeneratorExit:
    # close the file

---

如果你想要更简单的方法（不使用生成器的神秘 close 方法），可以让 cat 接受一个类似文件的对象，而不是一个字符串来打开，然后自己处理文件的输入输出：

for filename in filenames:
    with open( filename ) as theFile:
        for line in cat( theFile ):
            ...

---

不过，其实你基本上不需要担心这些，因为垃圾回收机制会处理所有这些事情。不过，

明确的做法总是比模糊的好

Answer 3

通过在同一个对象中实现上下文协议和迭代器协议，你可以写出非常棒的代码，比如下面这样：

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

这是一个示例实现，经过在Linux系统上使用Python 2.5测试。它会读取/etc/passwd文件中的每一行，直到找到用户audio的那一行，然后就停止了：

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())

或者更简单的写法：

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

文件对象实现了迭代器协议（详细信息见http://docs.python.org/library/stdtypes.html#file-objects）