Using 'yield' (i.e. generator) and 'return' together in a Python function

Question

I tried to create a simple Python function where the result is generator or normal output depending on a condition. But the function did not work. it seems the presence of the keyword 'yield' in the function leads the function to return a generator in all cases.

Below is a simple function that demonstrate what I tried to do.

# Generator function
def square_gen(numbers):
    if len(numbers) > 3:
        for number in numbers:
            yield number**2

    else:
        output= []
        for number in numbers:
            output.append(number**2)
        return output

The plan is to call the function in two different ways and get different corresponding outputs:

square_gen([1,2,3]) returns a generator object that I can iterate over (since length of the argument is not greater than 3) .
square_gen([1,2,3,4]) returns [1,4,9,16] (since length of the argument is greater than 3).

However, only 1) works. It returns a generator object that one can iterate over. Surprisingly, 2) still returns a generator object (I wonder why) but looping over this generator object results in an error.

So basically, I want to know why the non-generator branch of the function still returns a generator object.

Both the input should return a generator? Because the first case having len=3 should return a list — Will, Oct 30 '22 at 16:22
Does this answer your question? [Unexecuted yield statement blocks function to run?](https://stackoverflow.com/questions/42109265/unexecuted-yield-statement-blocks-function-to-run) — Ignatius Reilly, Oct 30 '22 at 16:22
You have already reached the correct conclusions about the behaviour of functions with `yield` in them. What are your clients going to do? Do they already know how long the parameter is that they supply? — quamrana, Oct 30 '22 at 16:22
From the docs: [Using yield in a function definition is sufficient to cause that definition to create a generator function instead of a normal function.](https://docs.python.org/3/reference/simple_stmts.html#the-yield-statement) — wwii, Oct 30 '22 at 16:55

score 4 · Answer 1 · answered Oct 30 '22 at 16:22

You are right in your assumption, having 'yield' anywhere in your method makes a generator, even if it's under an 'if'.

Not that there are many reasons to do it, but if you want to achieve what you asked you can do it like this

def square_gen(numbers):
    if len(numbers) > 3:

        def gen():
            for number in numbers:
                yield number**2
        return gen()

    else:
        output= []
        for number in numbers:
            output.append(number**2)
        return output

Using 'yield' (i.e. generator) and 'return' together in a Python function

1 Answers1