We have looked at list, set and dict comprehensions in previous articles.

We have seen how comprehensions work and how we can use them for filtering and transforming iterables to list, set or dict.

There is one thing left to discuss still.

We know that we can write any comprehension as a normal for loop.

Sometimes we need a for loop in a for loop, for example, when we have nested data structures such as a list of lists.

Below is a list that contains other lists, and they contain numbers:

list_of_lists = [
    [1, 2],
    [3, 4, 5],
    [6, 7, 8]
]

NestingPermalink

How would we turn the above list of lists into a simple list of numbers? We would need to iterate through each list in it and then through each number in each of those lists:

numbers = []
for sublist in list_of_lists:
    for n in sublist:
        numbers.append(n)

numbers
# outputs: [1, 2, 3, 4, 5, 6, 7, 8]

We can do it also with nested list comprehension:

numbers = [n for sublist in list_of_lists for n in sublist]
numbers
# outputs: [1, 2, 3, 4, 5, 6, 7, 8]

It might be a little confusing to read it all in one line the first time. But we can format our comprehensions any way we like it, for example like this:

numbers = [n
           for sublist in list_of_lists
           for n in sublist]

Now it looks much more similar to normal for loops.

The following image shows how nested comprehension maps to nested for loops:

We can nest as many for loops as we need, but most of the time, it's two of them; seldom, it is three. I don't think I've seen more in production code.

Filtering in nested comprehensionsPermalink

We can still filter what we want to include in our results with nested comprehensions, but we need to be careful where our if is.

Filtering if always comes right after the for part of the comprehension which we want to filter.

If we wanted only even numbers in our resulting list, we need to put if after the part of comprehension that is iterating through numbers in the sublist:

[n for sublist in list_of_lists for n in sublist if n % 2 == 0]
# outputs: [2, 4, 6, 8]

The comprehension above would translate to the following for loop:

numbers = []
for sublist in list_of_lists:
    for n in sublist:
        if n % 2 == 0:
            numbers.append(n)

The following image shows how the comprehension maps to the for loop.

So, when a comprehension has an if in it, it always belongs to the for on its left.

Say we want to use numbers only from lists with three numbers.

To do that, we would need to filter out some lists from our list_of_lists. We do it by putting if right after the for that iterates through list_of_lits:

[n for sublist in list_of_lists if len(sublist) == 3 for n in sublist]
# outputs: [3, 4, 5, 6, 7, 8]

Above comprehension would translate to the following for loop:

numbers = []
for sublist in list_of_lists:
    if len(sublist) == 3:
        for n in sublist:
            numbers.append(n)

The following image shows how the comprehension maps to the for loop.

During iteration at any level, the part after the if is only executed if condition in if evaluates to True - same as for normal for loops.

As you can guess, we can filter at all levels of comprehension if we need to.

If we wanted to process only lists that have three numbers in them and put only even numbers in our resulting list, this is how we should do it:

[n
 for sublist in list_of_lists if len(sublist) == 3
 for n in sublist if n % 2 == 0]

# outputs: [4, 6, 8]

We could do the same thing using a regular for loop:

numbers = []
for sublist in list_of_lists:
    if len(sublist) == 3:
        for n in sublist:
            if n % 2 == 0:
                numbers.append(n)

And here is how the comprehension maps to the for loop:

Trying to keep nested comprehension in one line can quickly turn ugly, so unless we can keep it short, I suggest splitting nested comprehension into multiple lines, especially if we also use filtering.

The following is still somewhat ok for an experienced developer:

numbers = [n for sublist in list_of_list for n in sublist]

But anything longer, I'd go for multiple lines:

squared = [n * n
           for sublist in list_of_lists if len(sublist) == 3
           for n in sublist if n % 2 == 0]

# OR even like this
squared = [n * n
           for sublist in list_of_lists 
           if len(sublist) == 3
           for n in sublist 
           if n % 2 == 0]

We can nest set and dict comprehension as well. Here's an example of building a dict from our list_of_lists where the key is a number and the value is its square. Also, with filtering for even numbers and only using lists with three items:

{n: n * n
 for sublist in list_of_lists if len(sublist) == 3
 for n in sublist if n % 2 == 0}

# outputs: {4: 16, 6: 36, 8: 64}

SummaryPermalink

• comprehensions can do multiple iterations traversing nested data
• we can filter at any level by adding if after the for part of each level
• it works same for list, set and dict comprehensions
• it helps to split nested comprehension into multiple lines

Happy coding!

You might also like