Eli Rose
Eli Rose

Reputation: 6988

Equivalent of `return` for Python generators

Sometimes, when rewriting recursive functions as generators, I miss the brevity of return.

"""
Returns a list of all length n strings that can be made out of a's and/or b's.
"""
def ab_star(n):
    if n == 0:
        return [""]

    results = []
    for s in ab_star(n - 1):
        results.append("a" + s)
        results.append("b" + s)

    return results

turns into

"""
Generator for all length n strings that can be made out of a's and/or b's.
"""
def ab_star(n):
    if n == 0:
        yield ""

    else:
        for s in ab_star(n - 1):
            yield "a" + s
            yield "b" + s

It's that else that bugs me. I wish there was a way to say "yield, and this is it, so exit the function". Is there a way?

Upvotes: 3

Views: 4790

Answers (2)

shx2
shx2

Reputation: 64298

Don't miss return, use it.

You can return right after you yield.

def ab_star(n):
    if n == 0:
        yield ""
        return
    for s in ab_star(n - 1):
        yield "a" + s
        yield "b" + s

An alternative is to use return in both cases, where the first case returns a sequence of length 1, and the second returns a generator-expression:

def ab_star(n):
    if n == 0:
        return ( "", )
    return ( c+s for s in ab_star(n - 1) for c in 'ab' )

This avoidance of yield avoids the limitation that you cannot use both return <value> and yield in the same function.

(This works in your case because your function doesn't have to be a generator. Since you only iterate over the results, it can also return a tuple.)

Upvotes: 7

Tim Peters
Tim Peters

Reputation: 70582

There isn't. When I wrote the "Simple Generators PEP", I noted:

Q. Then why not allow an expression on "return" too?

A. Perhaps we will someday.  In Icon, "return expr" means both "I'm
   done", and "but I have one final useful value to return too, and
   this is it".  At the start, and in the absence of compelling uses
   for "return expr", it's simply cleaner to use "yield" exclusively
   for delivering values.

But that never gained traction. Until it does ;-), you can make your generator look more like your first function by writing the first part as:

if n == 0:
    yield ""
    return

Then you can drop the else: statement and dedent the rest.

Upvotes: 5

Related Questions