Unpack List in List Comprehension

Question

listA = ["one", "two"]
listB = ["three"]
listC = ["four", "five", "six"]
listAll = listA + listB + listC
dictAll = {'all':listAll, 'A':listA, 'B':listB, 'C':listC,}


arg = ['foo', 'A', 'bar', 'B']
result = [dictAll[a] for a in arg if dictAll.has_key (a)]

I get the following result [['one', 'two'], ['three']] but what I want is ['one', 'two', 'three']

how do I unpack those lists in the list comprehension?

Answer 1

You can use itertools.chain.from_iterable :

>>> from itertools import chain
>>> list(chain.from_iterable(dictAll.get(a, []) for a in arg))
['one', 'two', 'three']

Also don't use dict.has_key it is deprecated(and removed in Python 3), you can simply check for a key using key in dict .

Answer 2

You can use a nested comprehension:

>>> [x for a in arg if dictAll.has_key(a) for x in dictAll[a]]
['one', 'two', 'three']

The order has always been confusing to me, but essentially it nests the same way it would if it were a loop. eg the left most iterable is the outermost loop and the right most iterable is the innermost loop.