How to include condition in composite list comprehension?

user_1142881 0
Viewed 56

I would like to flatten a list but keep NaNs. The following base code works when there are no NaNs:

l = [[1], [2, 3, 4], [5]]
[item for sublist in l for item in sublist]
> [1, 2, 3, 4, 5]

Now if I have the following case it will break:

import numpy as np
l = [[1], [2, 3, 4], np.nan, [5]]
[item for sublist in l for item in sublist]
> TypeError: 'float' object is not iterable

which makes sense, but I need this border case to be handled and I'm not sure how to add the special case condition in the list comprehension above. How can I modify such comprehension to cover such case?

The expected output is:

[1, 2, 3, 4, np.nan, 5]
python
2 Answers

I would check to see if the sublist is iterable. If not then wrap it in a tuple to then be flattened:

import numpy as np

l = [[1], [2, 3, 4], np.nan, [5]]

[item for sublist in l for item in (sublist if hasattr(sublist, "__iter__") else [sublist])]

>>> [1, 2, 3, 4, nan, 5]

Using chain from itertools would make it a bit cleaner:

from itertools import chain

list(chain(*(sublist if hasattr(sublist, "__iter__") else (sublist, ) for sublist in l)))
>>> [1, 2, 3, 4, nan, 5]

Source for checking if iterable (StackOverflow)

user_39360440

Trying different things I found a possible solution:

[item for sublist in l for item in (sublist if type(sublist)==list else [sublist])]
> [1, 2, 3, 4, np.nan, 5]
user_11428810
Related