动态扩充迭代器遍历异常修复咨询（需保留主程序结构）

遍历文件行时，当某行匹配特定模式需要优先遍历另一个文件，使用itertools.chain扩展迭代器后无法遍历新增内容，问题出在__iter__方法的for循环绑定了初始迭代器，后续更新不会生效。

原代码实现 #

import itertools

class LineGenerator():
    instance = None
    class __linegenerator():
        def __init__(self):
            self.generator = []; self.linecount = 0
        def inputfile(self,filename):
            print(f"add input file: {filename}")
            self.generator = itertools.chain( open(filename),self.generator )
        def __iter__(self):
            for l in self.generator:
                l = l.strip("\n")
                self.linecount += 1
                yield l
    def __new__(cls):
        if not LineGenerator.instance:
            LineGenerator.instance = LineGenerator.__linegenerator()
        return LineGenerator.instance
    def __getattr__(self,attr):
        return self.instance.__getattr__(attr)

if __name__=="__main__":
    with open("lines1","w") as lines1:
        lines1.write("a\n")
        lines1.write("b\n")
        lines1.write("c\n")
    with open("lines2","w") as lines2:
        lines2.write("aa\n")
        lines2.write("bb\n")
        lines2.write("cc\n")
    LineGenerator().inputfile("lines1")
    for l in LineGenerator():
        print(f"{l}")
        if l=="b":
            LineGenerator().inputfile("lines2")

当前输出 #

add input file: lines1
a
b
add input file: lines2
c

期望输出 #

add input file: lines1
a
b
add input file: lines2
aa
bb
cc
c

修复方案 #

修改__iter__方法，放弃直接for循环迭代器，改用手动调用next()的方式，确保每次迭代都使用当前最新的self.generator对象，主程序结构完全保留不变：

import itertools

class LineGenerator():
    instance = None
    class __linegenerator():
        def __init__(self):
            self.generator = []; self.linecount = 0
        def inputfile(self,filename):
            print(f"add input file: {filename}")
            self.generator = itertools.chain( open(filename),self.generator )
        def __iter__(self):
            while True:
                try:
                    # 每次从当前的generator中取元素
                    l = next(self.generator)
                    l = l.strip("\n")
                    self.linecount += 1
                    yield l
                except StopIteration:
                    # 迭代器耗尽时退出循环
                    break
    def __new__(cls):
        if not LineGenerator.instance:
            LineGenerator.instance = LineGenerator.__linegenerator()
        return LineGenerator.instance
    def __getattr__(self,attr):
        return self.instance.__getattr__(attr)

if __name__=="__main__":
    with open("lines1","w") as lines1:
        lines1.write("a\n")
        lines1.write("b\n")
        lines1.write("c\n")
    with open("lines2","w") as lines2:
        lines2.write("aa\n")
        lines2.write("bb\n")
        lines2.write("cc\n")
    LineGenerator().inputfile("lines1")
    for l in LineGenerator():
        print(f"{l}")
        if l=="b":
            LineGenerator().inputfile("lines2")

修复原理 #

原代码中for l in self.generator会在循环开始时将self.generator转换为迭代器并绑定，后续修改self.generator不会影响已经启动的循环。改用next()手动迭代后，每次循环都会从当前最新的self.generator中获取下一个元素，当调用inputfile替换self.generator为新的itertools.chain对象时，后续迭代会优先遍历新增的文件内容，再继续原有的迭代流程。

内容的提问来源于stack exchange，提问作者Victor Eijkhout