Ali SAID OMAR Ali SAID OMAR - 8 days ago 5
Python Question

Generator with return statement

During my coverage, I scratched my head on the following case (python 3.4)

def simple_gen_function(str_in, sep=""):
if sep == "":
yield str_in[0]
for c in str_in[1:]:
yield c
else:
return str_in
# yield from str_in

str_in = "je teste "
t = "".join(simple_gen_function(str_in))
p = "".join(simple_gen_function(str_in, "\n"))

print("%r %r" % (t, p))
# 'je teste' ''


Using return in the generator, the return was not "reached" while using
yield from str_in
I have the expected result.

The question seems simple, but I believed that using return in a generator, it was in reached.

Answer

The presence of yield in a function body turns it into a generator function instead of a normal function. And in a generator function, using return is a way of saying "The generator has ended, there are no more elements." By having the first statement of a generator method be return str_in, you are guaranteed to have a generator that returns no elements.

As a comment mentions, the return value is used as an argument to the StopIteration exception that gets raised when the generator has ended. See:

>>> gen = simple_gen_function("hello", "foo")
>>> next(gen)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration: hello

If there's a yield anywhere in your def, it's a generator!

In the comments, the asker mentions they thought the function turned into a generator dynamically, when the yield statement is executed. But this is not how it works! The decision is made before the code is ever excuted. If Python finds a yield anywhere at all under your def, it turns that def into a generator function.

See this ultra-condensed example:

>>> def foo():
...     if False:
...         yield "bar"
...     return "baz"
>>> foo()
<generator object foo at ...>
>>> # The return value "baz" is only exposed via StopIteration
>>> # You probably shouldn't use this behavior.
>>> next(foo())
Traceback (most recent call last):
  ...
StopIteration: baz
>>> # Nothing is ever yielded from the generator, so it generates no values.
>>> list(foo())
[]