Using a while loop walrus operator with a generator

1.2k views Asked by At

I'd like to use a walrus operator with a simple generator in a while loop

def gen():
    for i in range(5):
        yield i
    yield 10


g = gen()

while i := next(g):
    print(i)

I expect the output to be:

0
1
2
3
4
10

However it's not printing anything. I'm aware I can do this in a for loop:

for i in gen():
  print(i)

I'm not interested in the output being exact. I want a simple example of using generators with walrus operators.

3

There are 3 answers

1
testfile On

gen() is returning 0 on the first iteration. range(1, 5) to set range to go from 1 to 5. I also updated 10 to None else the generator runs out and throws Ex

def gen():
    for i in range(1, 5):
        yield i
    yield None


g = gen()

while i := next(g):
    print(i)

output:

1
2
3
4

Process finished with exit code 0
1
khelwood On

This loop:

while i := next(g):
    print(i)

will iterate through each element of your generator and print it, until one of them is falsey. Since your generator emits 0 as its first element, the while condition is immediately false, and the loop ends without printing.

But if, for example, your iterator was:

g = iter([2,1,0,2,5])

then your loop would print 2 and 1, and stop when it received the value 0.


If you want to iterate through all the elements until it runs out of elements, just use:

for i in g:
    print(i)
3
Freddy Mcloughlan On

Because you are using a while loop you need the condition to be True. The first element that is returned is 0. Which is interpreted as False.

If you need the := operator, use:

try:
    while (i := next(g)) or i == 0:
        print(i)
except StopIteration:
    # This executes when the generator is exhausted
    pass

Or use range(1, 10), so 0 is never returned


If you are open to changes, please use:

for i in g:
    print(i)

Or you avoid the removal of g's assignment, and do:

for i in gen():
    print(i)