CODEWITHSUNDEEP
pythonrecursiongeneratoryielditerable
ZheekChak
ZheekChak

Reputation: 53

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

Upvotes: 5

Views: 311

Answers (2)

Grismar
Grismar

Reputation: 31406

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)

Upvotes: 1

Paritosh Singh
Paritosh Singh

Reputation: 6246

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.

Upvotes: 6

Related Questions