CODEWITHSUNDEEP
pythongenerator
norok2
norok2

Reputation: 26896

Control recursion on nested lists / strings

Assume I have the following input:

items = [1, 2, [3, 4], (5, 6), 'ciao', range(3), (i for i in range(3, 6))]

and I want to perform some recursive operation on items.

For the sake of simplicity, let's say I want to flatten items (but could be anything else), one way of doing this would be:

def flatten(items, shallow=(str, bytes, bytearray)):
    for item in items:
        if isinstance(item, shallow):
            yield item
        else:
            try:
                yield from flatten(item)
            except TypeError:
                yield item

this would produce:

print(list(flatten(items)))
[1, 2, 3, 4, 5, 6, 'ciao', 0, 1, 2, 3, 4, 5]

Now how could I modify flatten() so that I could produce the following (for arbitrary nesting levels)?

print(list(flatten(items)))
[1, 2, 3, 4, 5, 6, 'c', 'i', 'a', 'o', 0, 1, 2, 3, 4, 5]

Upvotes: 2

Views: 69

Answers (1)

wim
wim

Reputation: 362945

Just add a length check next to the shallow check:

if isinstance(item, shallow) and len(item) == 1:

Upvotes: 5

Related Questions