Equivalent of `return` for Python generators

Question

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?

Answer 1

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.)