pythonyield

Surprising result with a conditional `yield`

I have the following Python code using yield:

def foo(arg):
    if arg:
        yield -1
    else:
        return range(5)

Specifically, the foo() method shall iterate over a single value (-1) if its argument is True and otherwise iterate over range(). But it doesn't:

>>> list(range(5))
[0, 1, 2, 3, 4]
>>> list(foo(True))
[-1]
>>> list(foo(False))
[]

For the last line, I would expect the same result as for the first line ([0, 1, 2, 3, 4]). Why is this not the case, and how should I change the code so that it works?

Solution

  • Using yield from seems to fix your function:

    import itertools

def foo(arg):
    if arg:
        yield -1
    else:
        yield from range(5)


print(list(foo(True)))
print(list(foo(False)))

    Output as requested