Why can't you toggle a function generator's behavior by an argument?

Question

Consider these two functions:

def foo():
    x = 0
    while True:
        yield x
        x += 1

def wrap_foo(limit=10, gen=True):
    fg = foo()
    count         


        

Answer 1

It seems like Python pre-classifies the function as a generator based on the presence of the yield statement

Yes, that's exactly what happens. wrap_foo is determined to be a generator at function definition time. You could consider using generator expressions instead:

def wrap_foo(limit=10, gen=True):
    fg = foo()
    if gen:
        return (next(fg) for _ in range(limit))
    else:
        return [next(fg) for _ in range(limit)]