Selective flattening of a Python list

Question

Suppose I have a list containing (among other things) sublists of different types:

[1, 2, [3, 4], {5, 6}]

that I'd like to flatten in a selective way, depending on the type of its elements (i.e. I'd like to only flatten sets, and leave the rest unflattened):

[1, 2, [3, 4], 5, 6]

My current solution is a function, but just for my intellectual curiosity, I wonder if it's possible to do it with a single list comprehension?

Answer 1

List comprehensions aren't designed for flattening (since they don't have a way to combine the values corresponding to multiple input items).

While you can get around this with nested list comprehensions, this requires each element in your top level list to be iterable.

Honestly, just use a function for this. It's the cleanest way.

Answer 2

Amber is probably right that a function is preferable for something like this. On the other hand, there's always room for a little variation. I'm assuming the nesting is never more than one level deep -- if it is ever more than one level deep, then you should definitely prefer a function for this. But if not, this is a potentially viable approach.

>>> from itertools import chain
>>> from collections import Set
>>> list(chain.from_iterable(x if isinstance(x, Set) else (x,) for x in l))
[1, 2, [3, 4], 5, 6]

The non-itertools way to do this would involve nested list comprehensions. Better to break that into two lines:

>>> packaged = (x if isinstance(x, collections.Set) else (x,) for x in l)
>>> [x for y in packaged for x in y]
[1, 2, [3, 4], 5, 6]

I don't have a strong intuition about whether either of these would be faster or slower than a straightforward function. These create lots of singleton tuples -- that's kind of a waste -- but they also happen at LC speed, which is usually pretty good.

Answer 3

You can use flatten function from funcy library:

from funcy import flatten, isa

flat_list = flatten(your_list, follow=isa(set))

You can also peek at its implementation.