Python: Function to flatten generator containing a

2019-04-21 16:30发布

问题:

I would like to know how to write python function which can flatten generator which yields another generators or iteables (which can also yield another generators/iterables ... possibly infinitely).

Here is example:

gen(gen(1,2,3), gen(4,5,6), [7,8,9], [gen(10,11,12), gen(13,14,15)])

note: gen - means generator object, content between parentheses after gen is data that will generator gen yield.

The expected result after "flattening": gen(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15)

It is necessary for the flatten function to return generator too! (because otherwise, the preceeding usage of generators would be meaningless).

Just to note, I am using python 3.

Thanks!

回答1:

The easiest way is a recursive flattening function. Assuming you want to descend into every iterable except for strings, you could do this:

def flatten(it):
    for x in it:
        if (isinstance(x, collections.Iterable) and
            not isinstance(x, str)):
            for y in flatten(x):
                yield y
        else:
            yield x

Starting from Python 3.3, you can also write

def flatten(it):
    for x in it:
        if (isinstance(x, collections.Iterable) and
            not isinstance(x, str)):
            yield from flatten(x)
        else:
            yield x


回答2:

The non-recursive method is essentially an unrolling of the recursive method, using a stack:

def flatten(it):
    stack = []
    it = iter(it)
    while True:
        try:
            x = next(it)
        except StopIteration:
            if stack:
                it = stack.pop()
                continue
            else:
                return
        if isinstance(x, collections.Iterable) and not isinstance(x, str):
            stack.append(it)
            it = iter(x)
        else:
            yield x