If else logic in python nested list comprehension

Question

I have a list which could consist of either lists or integers or None. I want to flatten the iterable elements (assuming only lists for now) inside this list to individual elements.

for eg:

[[0, 1], [2, 3], 1, 3, 4, 0, None] into [0,1,2,3,1,3,4,0,None] 

using list comprehension. I found another similar question but all those elements were iterable in that list, since mine has integers too how do I use if else logic in the list comprehension for the first "for loop". I am trying something like this, but not sure what the exact syntax is to flatten out.

[ item sublist if isinstance(sublist,list) else [sublist] for sublist in A for item in sublist ]

Based on other questions, if-else should occur prior to for loops and for loops have to occur in order. I am unable to insert if else after the first for loop, syntax allows only if and not else.

Could someone help with syntax of this please, for doing if-else on the first for loop or any intermediate for loop in nested for loops in comprehension?

Answer 1

You can use a generator to convert your mixed list into an iterable that only has lists:

gen = (x if isinstance(x, collections.Iterable) else [x] for x in A)

Then you can use the standard flattening idiom to flatten out the generator:

flattened = [y for x in gen for y in x]

Answer 2

@mgilson gave an elegant solution using a comprehension. It is also possible to do so in a natural loop using error-trapping:

items = [[0, 1], [2, 3], 1, 3, 4, 0, None]
flattened = []
for item in items:
    try:
        flattened.extend(item)
    except TypeError:
        flattened.append(item)

print(flattened) #prints [0, 1, 2, 3, 1, 3, 4, 0, None]