How does the magic flatten method work?

Question

I'm reading "Beginning python, from novice to professional", in which there is a magic flatten function which confused me.

def flatten(nested):
    try:
        for sublist in nested:
            for element in flatten(sublist):
                yield element
    except TypeError:
        yield nested

I know that yield returns an element. So assume I got a list L = [ 1, [[2]] ]. And I call it with this flatten(), like:

L = [ 1, [[2]] ]
for i in flatten(L):
    print i

1
2

I was really confused, when we call the for loop, won't we triger the flatten(), and we saw the first element 1, which no doubt will cause TypeError in the try block, shouldn't this returned nested variable in the except block be the whole list [1, [[2]] ]? why would it return 1 instead?

[What does the yield keyword do in Python?](https://stackoverflow.com/questions/231767/what-does-the-yield-keyword-do-in-python) — parchment, Oct 21 '14 at 13:21
See also https://docs.python.org/2/tutorial/classes.html#generators — jonrsharpe, Oct 21 '14 at 13:21
This function combines [generators](http://stackoverflow.com/q/231767/660921) and [recursion](http://stackoverflow.com/questions/717725/understanding-recursion). Two concepts that seem to confuse a lot of programmers. — Martin Tournoij, Oct 21 '14 at 13:21
Don't use this function with a list containing strings or you'll get infinite recursion (because strings are iterable, and their elements are also strings). — interjay, Oct 21 '14 at 13:22

score 1 · Answer 1 · answered Oct 21 '14 at 13:33

Perhaps rewriting the function without using yield can make it more understandable to you:

def flatten(nested):
    res = []
    try:
        for sublist in nested:
            for element in flatten(sublist):
                res.append(element)
    except TypeError:
        res.append(nested)
    return res

In the most basic terms, that is what yield does. It 'sends out' a value one at a time from a function. For a more complete explanation, see the answer to What does the yield keyword do in Python?.

score 1 · Answer 2 · answered Oct 21 '14 at 13:34

It can be helpful to throw in a print statement to see what's going on. Eg,

#! /usr/bin/env python

def flatten(nested, depth=0):
    print "%d: %r" % (depth, nested)
    try:
        for sublist in nested:
            for element in flatten(sublist, depth+1):
                yield element
    except TypeError:
        yield nested


def main():
    L = [ 1, [[2]] ]
    for i in flatten(L):
        print "output %s\n" % i


if __name__ == '__main__':
    main()

output

0: [1, [[2]]]
1: 1
output 1

1: [[2]]
2: [2]
3: 2
output 2

score 0 · Answer 3 · answered Oct 22 '14 at 08:20

I think I understood what happened here now.

def flatten(nested):
    try:
        for sublist in nested:
            for element in flatten(sublist):
                yield element
    except TypeError:
        yield nested

We still have L = [1, [[2]]], and we call flatten(L).

Now the first element is 1, now we have moved to for element in flatten(1).

As when we call flatten(), and meet TypeError, we'll get nested variable, at here, it is 1.

So we got 1.

And the original L now have only [[2]] left in it.

As we called a for loop: for i in flatten(L), things will goes on.

We will now move to for element in flatten([[2]]), yes, we got an yield element below it, but we have to execute for element in flatten([[2]]) first and saw what happened.

Now we'll move to for element in flatten([2]), it's a temptation to go straight to the yield element below it. But hold on, let's go into flatten another time.

Now we have move to flatten(2), it will raise TypeError like flatten(1) do.

So we got 2 this time.

Things explained!

How does the magic flatten method work?

3 Answers3

Linked

Related