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

Question

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

Lithopone

Answer 1

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.

Incertitude

Answer 2

This fixes your code:

from collections.abc import Iterable


def flatten(nested):
    for item in nested:
        if isinstance(item, Iterable):
            yield from flatten(item)
        else:
            yield item


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

for element in y:
    print(element)

This works because you call flatten again, but forget to yield from (after all, the new call also returns a generator)

Note that isinstance(item, Iterable) is probably not what you want, since it will break for strings. Strings are Iterable, but in a for loop, the characters that are returned from it are strings themselves. Probably better to check if it is a list.

def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item


for element in flatten([1, 2, ['three', 4, [5, [[6, 7]]]]]):
    print(element)

Fardel

You just reminded me a tricky point in python : characters are stings! So if I include a string in my list, it will flatten it into characters, however characters are still stings so it keeps flattening.....That's a disaster. Thanks for noting me of that! – Lithopone