Optional yield or return in python3. How to?

Question

I would like to have a function that can, optionally, return or yield the result. Here is an example.

def f(option=True):
    ...
    for...:
        if option:
            yield result
        else:
            results.append(result)

    if not option:
        return results

Of course, this doesn't work, I have tried with python3 and I always get a generator no matter what option value I set. As far I have understood, python checks the body of the function and if a yield is present, then the result will be a generator. Is there any way to get around this and make a function that can return or yield at will?

Answer 1

You can't. Any use of yield makes the function a generator.

You could wrap your function with one that uses list() to store all values the generator produces in a list object and returns that:

def f_wrapper(option=True):
    gen = f()
    if option:
        return gen    # return the generator unchanged
    return list(gen)  # return all values of the generator as a list

However, generally speaking, this is bad design. Don't have your functions alter behaviour like this; stick to one return type (a generator or an object) and don't have it switch between the two.

Consider splitting this into two functions instead:

def f():
    yield result

def f_as_list():
    return list(f())

and use either f() if you need the generator, and f_as_list() if you want to have a list instead.

Since list(), (and next() to access just one value of a generator) are built-in functions, you rarely need to use a wrapper. Just call those functions directly:

# access elements one by one
gen = f()
one_value = next(gen)

# convert the generator to a list
all_values = list(f())

Answer 2

What about this?

def make_f_or_generator(option):
    def f():
        return "I am a function."
    def g():
        yield "I am a generator."
    if option:
        return f
    else:
        return g

This gives you at least the choice to create a function or a generator.