Issue with a python function returning a generator or a normal object

Question

I defined the function f as

def f(flag):
    n = 10
    if flag:
        for i in range(n):
            yield i
    else:
        return range(n)

But f returns a generator no matter what flag is:

>>> f(True)
<generator object f at 0x0000000003C5EEA0>

>>> f(False)
<generator object f at 0x0000000007AC4828>

And if I iterate over the returned object:

# prints normally
for i in f(True):
    print(i)

# doesn't print
for i in f(False):
    print(i)

It looks like f(False) returns a generator which has been iterated over. What's the reason? Thank you.

Answer 1

A function containing a yield statement always returns a generator object.

Only when you iterate over that generator object will the code in the function be executed. Until that time, no code in the function is executed and Python cannot know that you'll just return.

Note that using return in a generator function has different semantics than in a regular function; return in this case simply is seen as 'exit the generator here'; the return value is discarded as a generator can only produce values via yield expressions.

It looks like you want to use yield from instead:

def f(flag):
    n = 10
    if flag:
        for i in range(n):
            yield i
    else:
        yield from range(n)

yield from requires Python 3.3 or up.

See the yield expression documentation:

Using a yield expression in a function’s body causes that function to be a generator.

When a generator function is called, it returns an iterator known as a generator. That generator then controls the execution of a generator function. The execution starts when one of the generator’s methods is called. At that time, the execution proceeds to the first yield expression, where it is suspended again, returning the value of expression_list to the generator’s caller.

Iteration over a generator calls the generator.__next__() method, triggering execution.

If you wanted to return a generator some of the time, then don't use yield in this function. You'd produce the generator by other means; using a separate function for example, or by using a generator expression perhaps:

def f(flag):
    n = 10
    if flag:
        return (i for i in range(n))
    else:
        return range(n)

Now no yield is used in f and it will no longer produce a generator object directly. Instead, the generator expression (i for i in range(n)) produces it, but only conditionally.