Search code examples
pythonpython-itertools

iterating over dynamically augmented iterator


I am iterating over lines of a file, but if a line matches a certain pattern, I first need to recursively iterate over another file. I'm using itertools.chain to extend the iterator. The problem is (probably) that the for loop in the __iter__ method makes a copy of the initial iterator so I'm not iterating over the updated one.

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

(I hope you can forgive my use of a singleton pattern)

Output:

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

Desired output:

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

Question: how can this design be fixed? I would like to keep the main program as it stands.


Solution

  • A way that achieves the result (changed a bit, as discussed in the comments). It keeps iterators on a stack and always gives the next line of the top iterator:

    class LineIterator:
    
        def __init__(self):
            self.iterator_stack = []
    
        def __iter__(self):
            return self
    
        def insertfile(self, filename):
            print(f"add input file: {filename}")
            def gen():
                with open(filename) as f:
                    yield from f
            self.iterator_stack.append(gen())
    
        def __next__(self):
            stack = self.iterator_stack
            while stack:
                try:
                    return next(stack[-1]).strip("\n")
                except StopIteration:
                    stack.pop()
            raise StopIteration
    
    
    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")
        lines = LineIterator()
        lines.insertfile("lines1")
        for l in lines:
            print(f"{l}")
            if l=="b":
                lines.insertfile("lines2")
    

    Output (Try it online!):

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