Search code examples
pythoniteratorstopiteration

Return from an iterator and then throw StopIteration


What would be the nice way to return something from an iterator one last time when it's exhausted. I'm using a flag, but this is rather ugly:

class Example():

    def __iter__(self):
        self.lst = [1,2,3]
        self.stop = False # <-- ugly            
        return self

    def next(self):
        if self.stop:  # <-- ugly
            raise StopIteration
        if len(self.lst) == 0:
            self.stop = True            
            return "one last time"
        return self.lst.pop()

Background: I'm fetching an unknown amount of strings from an external source and send them further down to the caller. When the process is over, I want to emit a string "x records processed". I have no control over calling code, so this must be done inside my iterator.


Solution

  • You could just yield from __iter__ which would turn it into a generator function (alternately you could just write a generator function as suggested by Dan). Just as a warning, this might be misleading to people that abuse the next method.

    class Example():
    
        def __iter__(self):
            lst = [1,2,3]
            for i in reversed(lst):
                yield i
            yield "one last time"