检查并关闭 python 生成器

Question

我有一段代码使用生成器从两个大文件中读取，并在两个文件之一中达到 EOF 时停止。我想知道（1）哪个生成器首先到达 EOF，（2）每个生成器的进度，即当第一个生成器到达 EOF 时，生成器中 i 的值（参见下面的代码），以及（3）另一个生成器中剩余的行数。我提前不知道每个文件有多长，并希望避免预先扫描文件。

我知道我可以通过以下方式获得进度：

1. 每次调用 next() 时增加一个计数器（这很丑！），或者

2. 让生成器返回一个计数器（见代码中的counter1 和counter2），

但在这两种情况下，我都不知道gen1 或gen2 中的哪一个达到了EOF。

我还发现我可以在StopIteration 异常中添加“消息”，但我想知道是否有更好的方法。在第一个 try...except 块之后，我能以某种方式找出哪个尚未达到 EOF 并推进它吗？ （我尝试在生成器上使用close() 或throw()，或者在生成器中使用finally 子句，但并没有真正理解它们。）

def gen1(fp):
    for i, line in enumerate(fp):
        int_val = process_line(line)
        yield int_val, i
    raise StopIteration, ("gen1", i)

def gen2(fp):
    for i, line in enumerate(fp):
        float_val = process_line_some_other_way(line)
        yield float_val, i
    raise StopIteration, ("gen2", i)

g1 = gen1(open('large_file', 'r'))
g2 = gen2(open('another_large_file', 'r'))

try:
    val1, counter1 = next(g1)
    val2, counter2 = next(g2)
    progress += 1
    while True:  # actual code is a bit more complicated than shown here
        while val1 > val2:
            val2, counter2 = next(g2)
        while val1 < val2:
            val1, counter1 = next(g1)
        if val1 == val2:
            do_something()
            val1, counter1 = next(g1)
            val2, counter2 = next(g2)

except StopIteration as err:
    first_gen_name, first_num_lines = err.args

gen1_finished_first = gen_name == 'gen1'

# Go through the rest of the other generator to get the total number of lines
the_remaining_generator = g2 if gen1_finished_first else g1
try:
    while True:
        next(the_remaining_generator)
except StopIteration as err:
    second_gen_name, second_num_lines = err.args

if gen1_finished_first:
    print 'gen1 finished first, it had {} lines.'.format(first_num_lines) # same as `counter1`
    print 'gen2 was at line {} when gen1 finished.'.format(counter2)
    print 'gen2 had {} lines total.'.format(second_num_lines)
else:
    ... # omitted

Answer 1

我认为您可能想改用an iterator class —— 它是使用标准 Python 类实现的，并且可以具有您需要的任何额外属性（例如 exhausted 标志）。

类似于以下内容：

# untested
class file_iter():
    def __init__(self, file_name):
        self.file = open(file_name)
        self.counted_lines = 0
        self.exhausted = False
    def __iter__(self):
        return self
    def __next__(self):
        if self.exhausted:
            raise StopIteration
        try:
            next_line = self.file.readline()
            self.counted_lines += 1
            return next_line
        except EOFError:
            self.file.close()
            self.exhausted = True
            raise StopIteration

Answer 2

您可以使用chain 在生成器末尾添加一个特殊的EOF 值。例如。

from itertools import chain
EOF = object()
fin = open('somefile')
src = enumerate(chain(fin, [EOF]))
while True:
    idx, row = next(src)
    if row == EOF:
        break  # End of file
    print idx, row

您也可以使用izip_longest。用你的生成器替换 f1 和 f2

from itertools import count, izip_longest
EOF = object()
with open('f1') as f1, open('f2') as f2:
    for i, r1, r2 in izip_longest(count(), f1, f2, fillvalue=EOF):
        if EOF in (r1, r2):
            print i, r1, r2
            break