Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How does the 'yield' keyword in python really work, especially when it comes with recursion?

I'm using python to flatten a nested list, like [1,2,[3,4,[5,[[6,7]]]]], I want to create an generator so I can use for loop to print all the numbers one by one in the nested list. But it just does not work as I expected.

When I replace the 'yield' keyword with 'print', the numbers get printed one by one. However in that way it is no longer a generator.

The following does not work properly:

from collections.abc import Iterable

def flatten(nested):
    for item in nested:
        if isinstance(item, Iterable):
            flatten(item)
        else:
            yield item
x = [1,2,[3,4,[5,[[6,7]]]]]
y = flatten(x)

for element in y:
    print(element)

However, if I wrote the code like below, where I replaced the yield with print, the numbers will get printed properly

from collections.abc import Iterable

def flatten(nested):
    for item in nested:
        if isinstance(item, Iterable):
            flatten(item)
        else:
            print(item)
x = [1,2,[3,4,[5,[[6,7]]]]]
y = flatten(x)

With the 'yield', if I wrote:

x = [[1,2],[3,4,[5,[[6,7]]]]]
y = flatten(x)
y.__next__()

The error message will be y.__next__() StopIteration

like image 811
ZheekChak Avatar asked Mar 04 '23 15:03

ZheekChak


1 Answers

You're never returning or yielding from the recursive calls. add a yield from and it should work.

from collections.abc import Iterable

def flatten(nested):
    for item in nested:
        if isinstance(item, Iterable):
            yield from flatten(item) #change here.
        else:
            yield item
x = [1,2,[3,4,[5,[[6,7]]]]]
y = flatten(x)

for element in y:
    print(element)
#Output:
1
2
3
4
5
6
7

Note that this problem was also inherently present in your original function, whether you would have used a yield or a return. This is why you should be careful when using print only and no returns in recursive calls, it can mask the fact that while the code runs properly, the outputs are not correctly captured/used.

like image 152
Paritosh Singh Avatar answered Apr 28 '23 04:04

Paritosh Singh