HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Recursive flatten with lazy evaluation iterator

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
evaluationiteratorwithrecursivelazyflatten

Problem

I'm trying to rewrite this:

def flatten(lst):
    flat = []
    for x in lst:
        if hasattr(x, '__iter__') and not isinstance(x, basestring):
            flat.extend(flatten(x))
        else:
            flat.append(x)
    return flat

In [21]: a=[1, [2, 3, 4, [5, 6]], 7]

In [22]: flatten(a)
Out[22]: [1, 2, 3, 4, 5, 6, 7]


..into a version that would flatten the iterable, but return values in a lazy manner. Now, this works:

def flat_fugly(s):
    if iterable(s):
        for x in s:
            yield chain.from_iterable(flat_fugly(x))
    else:
        yield takewhile(lambda x: True, [s])

list(islice(chain.from_iterable(flat_fugly(a)), 0, 6))
Out[34]: [1, 2, 3, 4, 5, 6]


But, as the name suggests... Is there a cleaner/better way?

Not to mention that I have to apply chain.from_iterable to flat_fugly anyway while I'd prefer to have plain iterator (I could wrap it in another function that would use chain.from_iterable of course, but still if it could be all made to fit in one more elegant function that would be preferable).

Solution

Not my idea, but I think this is better:

import collections

def flatten(lst):
    for item in lst:
        if isinstance(item, collections.Iterable) and not isinstance(item, basestring):
            for sublst in flatten(item):
                yield sublst
        else:
            yield item

Code Snippets

import collections


def flatten(lst):
    for item in lst:
        if isinstance(item, collections.Iterable) and not isinstance(item, basestring):
            for sublst in flatten(item):
                yield sublst
        else:
            yield item

Context

StackExchange Code Review Q#49877, answer score: 3

Revisions (0)

No revisions yet.