Nested list comprehensions

Are nested list comprehensions a bad idea?

Nested list comprehensions

Here's a nested list comprehension:

>>> scores_by_group = [[85, 92, 78], [91, 88, 95], [77, 82, 90]]
>>> percentages = [[score/100 for score in group_scores] for group_scores in scores_by_group]

This is a list comprehension with another list comprehension inside its mapping part.

List comprehensions create a new list, and since we've embedded a comprehension within a comprehension, this code creates a list-of-lists:

>>> percentages
[[0.85, 0.92, 0.78], [0.91, 0.88, 0.95], [0.77, 0.82, 0.9]]

I don't find this code very readable:

>>> percentages = [[score/100 for score in group_scores] for group_scores in scores_by_group]

But I don't think the comprehensions are the problem here. The problem is our use of whitespace.

Here's the same code broken up over multiple lines:

>>> percentages = [
...     [score/100 for score in group_scores]
...     for group_scores in scores_by_group
... ]
>>> percentages
[[0.85, 0.92, 0.78], [0.91, 0.88, 0.95], [0.77, 0.82, 0.9]]

I find this code much more readable.

But is it more or less readable than an equivalent for loop would be?

Nested for loops

Here's the same code without any list comprehensions:

>>> percentages = []
>>> for group_scores in scores_by_group:
...     group_percentages = []
...     for score in group_scores:
...         group_percentages.append(score/100)
...     percentages.append(group_percentages)
...
>>> percentages
[[0.85, 0.92, 0.78], [0.91, 0.88, 0.95], [0.77, 0.82, 0.9]]

We have a loop within a loop here.

The inner loop builds up a new list for each iteration of the outer loop, and the outer loop appends each of those inner lists to an outer list.

This for loop is pretty readable, but the purpose of this code might not be obvious at a glance.

Here's the same code with an inner comprehension and an outer loop:

>>> percentages = []
>>> for group_scores in scores_by_group:
...     percentages.append([score/100 for score in group_scores])
...
>>> percentages
[[0.85, 0.92, 0.78], [0.91, 0.88, 0.95], [0.77, 0.82, 0.9]]

To me, the purpose of this code seems a bit more obvious at a glance than the loop within a loop approach.

Comprehensions versus for loops

Python's for loops are for looping over an iterable while doing something.

So when I see a loop within a loop, I think "nested looping". I don't immediately think "we're building up a list-of-lists". I'd need to read through most of this code to figure out its purpose.

Comprehensions are specifically for building up new lists.

So whenever I see a comprehension, I think "we're building a new list". And when I see nested comprehensions, I think "we're building a list-of-lists":

>>> percentages = [
...     [score/100 for score in group_scores]
...     for group_scores in scores_by_group
... ]

Python's for loops showcase the looping, while comprehensions showcase the list-making.

When we use comprehensions our code looks like the data structure that we're trying to create. That's not the case when we use for loops.

Overly complex comprehensions

When I'm building up a new list, I usually prefer to use a comprehension, if I reasonably can.

Note that I said "reasonably". When it comes to using comprehensions, what is reasonable?

For example, is this reasonable?

students_by_group = [
    [{'name': 'Alice', 'score': 85}, {'name': 'Bob', 'score': 92}],
    [{'name': 'Charlie', 'score': 78}, {'name': 'Diana', 'score': 91}]
]
results = [
    [
        f"{s['name']}: {'Pass' if s['score'] >= 80 else 'Fail'} ({s['score']/100:.0%})"
        for s in group
    ]
    for group in students_by_group
    if len(group) > 0
]

This is a comprehension within a comprehension, and both the inner and outer comprehensions have a lot of logic within them.

The result is a list within a list, just as before:

>>> results
[['Alice: Pass (85%)', 'Bob: Pass (92%)'], ['Charlie: Fail (78%)', 'Diana: Pass (91%)']]

But this code isn't very readable.

Interestingly, unwrapping those into a for loop also isn't very readable:

results = []
for group in students_by_group:
    if len(group) > 0:
        group_results = []
        for student in group:
            status = 'Pass' if student['score'] >= 80 else 'Fail'
            result = f"{student['name']}: {status} ({student['score']/100:.0%})"
            group_results.append(result)
        results.append(group_results)

We could make this code a bit easier to read by creating a function or two to help give names to some of the operations we're performing.

So here we're using this score_status function within our append call in our inner loop:

def score_status(student):
    status = 'Pass' if student['score'] >= 80 else 'Fail'
    return f"{student['name']}: {status} ({student['score']/100:.0%})"


results = []
for group in students:
    if group:
        group_results = []
        for student in group:
            group_results.append(score_status(student))
        results.append(group_results)

This for loop is a lot more readable than the comprehension that we started with.

But now that we've broken up our code, would a comprehension be more readable in this case?

results = [
    [score_status(s) for s in group]
    for group in students
    if group
]

That comprehension doesn't look so bad.

Sometimes an unreadable comprehension may be a hint that your code has too much logic shoved into one loop.

Flattening lists with nested loops

Sometimes the phrase "nested comprehension" is used to describe a different sort of nesting.

We just looked at a comprehension inside of a comprehension, but it's also possible to write comprehensions that contain multiple loops.

Here's a comprehension with nested loops:

>>> blog_posts = [
...     {'title': 'Python Tips', 'tags': ['python', 'programming', 'tips']},
...     {'title': 'Web Dev Guide', 'tags': ['web', 'javascript', 'html']},
...     {'title': 'Data Science 101', 'tags': ['python', 'data', 'science']}
... ]
>>> [tag for post in blog_posts for tag in post['tags']]
['python', 'programming', 'tips', 'web', 'javascript', 'html', 'python', 'data', 'science']

That comprehension is equivalent to these for loops:

>>> all_tags = []
>>> for post in blog_posts:
...     for tag in post['tags']:
...         all_tags.append(tag)
...
>>> all_tags
['python', 'programming', 'tips', 'web', 'javascript', 'html', 'python', 'data', 'science']

Some Python programmers always avoid writing comprehensions with multiple loops. I don't agree with that approach.

While I do agree that this comprehension is less readable than that for loop, again, it's due to the lack of whitespace.

Breaking that comprehension up over multiple lines makes it easier to see what's happening:

>>> [
...     tag
...     for post in blog_posts
...     for tag in post['tags']
... ]
['python', 'programming', 'tips', 'web', 'javascript', 'html', 'python', 'data', 'science']

Reading order in multi-loop comprehensions

Keep in mind that because the loops in a comprehension must be in the same order as equivalent for loops, it does read a little funny.

We start with the outer loop and then work downward, which means the variable that we use in that first mapping section of our comprehension often isn't introduced until the last for line:

>>> [
...     tag
...     for post in blog_posts
...     for tag in post['tags']
... ]

Python often reads like English, but it really doesn't in the case of comprehensions with nested loops.

Despite the order of the for clauses sounding a bit awkward in English, I do sometimes use comprehensions with nested loops.

Whenever I see nested loops in a comprehension, I think "flattening". We're taking an iterable-of-iterables and flattening it into a single list.

But when I see nested for loops, I don't necessarily think "flattening". After all, our first example also involved nested for loops, but it didn't flatten anything:

>>> percentages = []
>>> for group_scores in scores_by_group:
...     group_percentages = []
...     for score in group_scores:
...         group_percentages.append(score/100)
...     percentages.append(group_percentages)
...
>>> percentages
[[0.85, 0.92, 0.78], [0.91, 0.88, 0.95], [0.77, 0.82, 0.9]]

Just as with single-loop comprehensions, multi-loop comprehensions can be a mental shortcut that helps us quickly identify the purpose of some code, specifically to flatten nested iterables.

Nested comprehensions can be pretty readable

Whether we're talking about comprehensions with multiple loops, or comprehensions with comprehensions within them, nested comprehensions can sometimes be helpful tools for making more readable Python code.

But readability is a subjective measure.

So the next time you find yourself making a list of lists or flattening a list of lists, try out the comprehension approach. And then ask yourself, "is the for loop or the comprehension more readable?"