CODEWITHSUNDEEP
pythonpython-3.x
urschrei
urschrei

Reputation: 26879

How can I wrap a function returning successive values in a generator?

I have a function bar(obj) which returns successive values from an object until it's exhausted, at which point it returns False (ignore the fact that it's not a method on Foo, or that I'm returning integers – this is just an MVCE to demonstrate the consumption of something). I'd like to wrap that function in a generator so I can lazily iterate over the values, but I'm clearly missing some logic, because this doesn't work:

class Foo:
    def __init__(self):
        self.counter = 0

    def inc(self):
        if self.counter <= 9:
            self.counter += 1
            return self.counter
        else:
            return False


def bar(obj):
    return obj.inc()


def iterbar(obj):
    res = bar(obj)
    if res:
        yield res
    else:
        raise StopIteration

foo = Foo()
lazy = iterbar(foo)
next(lazy) # this yields 1, as expected
next(lazy) # this immediately raises StopIteration from somewhere other than iterbar

How can I modify iterbar to work correctly?

Upvotes: 1

Views: 77

Answers (1)

Carcigenicate
Carcigenicate

Reputation: 45806

You aren't looping to repeat the calls to bar. Once you call next a second time, it picks up where the yield left off, gets to the end of the function, then indicates that by raising.

Just loop, then return to stop:

def iterbar(obj):
    while True:
        res = bar(obj)
        if res:
            yield res
        else:
            return

Or, if you want to be fancy and have Python 3.8+, you can use an assignment expression to make this more terse:

def iterbar(obj):
    while res := bar(obj):
        yield res

Upvotes: 2

Related Questions