CODEWITHSUNDEEP
pythoniterationgeneratoryield
Ziyuan
Ziyuan

Reputation: 4558

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

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.

Upvotes: 15

Views: 8826

Answers (3)

Denrey Christy
Denrey Christy

Reputation: 19

The generator doesn't just automatically produce the next yield value, unless you specifically call something like

next()

on it. I had a generator named "generate parameters", and it worked correctly once I did:

print(next(generate_parameters()))

Upvotes: 0

dano
dano

Reputation: 94881

You can work around this by using a nested function that actually uses yield:

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

>>> f(True)
<generator object gen at 0x7f62017e3730>
>>> f(False)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

As Martijn points out, any function containing a yield will always return a generator object, so if in some cases you want the body of f to actually be executed when f() is called, rather than only be executed when iterated over, you have to use this approach.

The standard library instance method map from concurrent.Futures.ProcessPoolExecutor/concurrent.Futures.ThreadPoolExecutor uses this to ensure futures are submitted as soon as map is called, rather than only when you try to actually retrieve results from it, for example:

def map(self, fn, *iterables, timeout=None):
    if timeout is not None:
        end_time = timeout + time.time()

    fs = [self.submit(fn, *args) for args in zip(*iterables)]

    # Yield must be hidden in closure so that the futures are submitted
    # before the first iterator value is required.
    def result_iterator():
        try:
            for future in fs:
                if timeout is None:
                    yield future.result()
                else:
                    yield future.result(end_time - time.time())
        finally:
            for future in fs:
                future.cancel()
    return result_iterator()

Upvotes: 6

Martijn Pieters
Martijn Pieters

Reputation: 1121924

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.

Upvotes: 24

Related Questions