Turning a for loop into a list comprehension

Let's turn a for loop into a list comprehension.

A for loop that builds up a new list

We have a list of strings that represent Python Morsels screencast names:

screencasts = [
    "Data structures contain pointers",
    "What is self?",
    "What is a class?",
    "Slicing",
    "How to make a function",
    "Methods are just functions attached to classes",
]

And we have a for loop that loops over this screencasts list and makes a new list (titles) of strings, changing each of the strings a little bit along the way:

titles = []
for name in screencasts:
    titles.append(name.title())

We've title-cased these strings, so the first letter of each word is capitalized.

>>> from pprint import pprint
>>> pprint(titles)
['Data Structures Contain Pointers',
 'What Is Self?',
 'What Is A Class?',
 'Slicing',
 'How To Make A Function',
 'Methods Are Just Functions Attached To Classes']

Copy-pasting our way into a comprehension

Whenever you have a for loop that builds up a new list and it's written in this format:

new_list = []
for item in some_list:
    new_list.append(some_operation_with(item))

You can copy-paste that for loop into a comprehension.

Let's copy-paste our for loop into a comprehension:

titles = []
for name in screencasts:
    titles.append(name.title())

We'll start our copy-pasting with the square brackets ([]):

titles = [
]

This square brackets indicate that we're building up a list.

Then we'll copy the thing that we're appending (name.title() in our case):

titles = [
    name.title()
]

Our looping logic (the for line without the colon) goes after that:

titles = [
    name.title()
    for name in screencasts
]

That comprehension is equivalent to the for loop we started with:

titles = []
for name in screencasts:
    titles.append(name.title())

Notice that we wrote our comprehension over multiple lines of code. We could have written the comprehension all on one line of code, but it wouldn't have been as readable.

titles = [name.title() for name in screencasts]

Whitespace is your friend, especially when it comes to code readability. I personally prefer to write every component of my comprehension on a separate line.

Making comprehensions that filter items

This strategy of copy-pasting loops into comprehensions also works when we have to filter elements down (only including items that match a specific condition).

This for loop builds up a new list, appending names that are 30 characters or greater:

long_names = []
for name in titles:
    if len(name) > 30:
        long_names.append(name[:27] + "...")

The resulting list (long_names) only includes long strings:

>>> long_names
['Data Structures Contain Poi...', 'Methods Are Just Functions ...']

Whenever you encounter a for loop written in this format:

new_list = []
for item in some_list:
    if condition_on(item):
        new_list.append(some_operation_with(item))

You can copy-paste your way into a comprehension.

Just as before, the first thing we'll copy is the square brackets ([]):

long_titles = [
]

The first thing we put after the opening square bracket is the thing that we're appending to our new list:

long_titles = [
    name.title()
]

Our looping logic (the for line) comes next:

long_titles = [
    name[:27] + "..."
    for name in names
]

And the last thing that we copy paste is our condition (the if line):

long_titles = [
    name[:27] + "..."
    for name in titles
    if len(name) > 30
]

This comprehension is written over multiple lines of code because I find it more readable than if we'd written the same logic squished all on to one line of code:

long_names = [name[:27] + "..." for name in titles if len(name) > 30]

Breaking each component of our comprehension onto its own line is a bit more readable:

long_titles = [
    name[:27] + "..."
    for name in titles
    if len(name) > 30
]

Turning nested for loops into a comprehension

This strategy of copy-pasting a for loop into comprehension even works in more complex cases.

For example, nested for loops can be turned into comprehensions:

uppercase_letters = []
for name in screencasts:
    for char in name:
        if char.isupper():
            uppercase_letters.append(char)

you can copy-paste your way into a comprehension in a nested case as well.

As long as your for loop only consists of for lines, if lines, and a single append line, you can copy paste your way into a comprehension.

The first thing we copy is (again) the square brackets:

uppercase_letters = [
]

The first thing we put inside our square brackets is the thing we're appending:

uppercase_letters = [
    char
]

After that come our looping logic (the for lines) and our condition (the if line) but with no colons at the end of the lines:

uppercase_letters = [
    char
    for name in screencasts
    for char in name
    if char.isupper()
]

Note that while the thing we're appending comes first, all the for and if lines remain in the same order as they were in our for loop.

What if directly copy-pasting isn't possible?

This copy-pasting strategy only works if your for loop is written in the correct format. If you call append multiple times or you have a bit more complex logic in your for loop (something that's not just if or for) then you're out of luck.

For example this for loop below has multiple append calls and it has an else block:

shortened = []
for name in screencasts:
    if len(name) > 30:
        shortened.append(name[:30-3] + "...")
    else:
        shortened.append(name)

In this particular case, we could collapse our conditional logic into a single append call if we make a helper function:

def ellipsify(name):
    if len(name) > 30:
        return name[:30-3] + "..."
    else:
        return name

This helper function contains the same logic that was in our for loop's if-else block.

Since that helper function is taking care of the if-else logic and returning a single result, we can rewrite our for loop to appending that one result to our new list:

shortened = []
for name in screencasts:
    shortened.append(ellipsify(name))

That for loop is now written in the correct format for copy-pasting our way into a comprehension:

shortened = [
    ellipsify(name)
    for name in screencasts
]

Note that in this particular case, an inline if statement could be used if you really wanted to avoid a factory function, but I don't recommend it.

Summary

If you have a for loop written in this format:

new_list = []
for item in some_list:
    new_list.append(some_operation_with(item))

You can copy-paste your way into an equivalent comprehension:

new_list = [
    some_operation_with(item)
    for item in some_list
]

If your for loop filters down based on a condition:

new_list = []
for item in some_list:
    if condition_on(item):
        new_list.append(some_operation_with(item))

You can also copy-paste your way into an equivalent comprehension:

new_list = [
    some_operation_with(item)
    for item in some_list
    if condition_on(item)
]

Your loop can even have multiple for or if statements:

new_list = []
for some_iterable in some_list:
    for item in some_iterable:
         if condition_on(item):
            new_list.append(some_operation_with(item))

and it only has a single append call, you can copy-paste your way from a for loop into a list comprehension.

You'll always copy-paste the thing you're appending first followed by the for and if lines in the same order they appear within your loop:

new_list = [
    some_operation_with(item)
    for some_iterable in some_list
    for item in some_iterable
     if condition_on(item)
]

It doesn't matter how many for lines you have or how many if lines you have, as long as you only have a single append call.

If your for loop is in a different format (you have more complex logic within it), you'll only be able to copy-paste it into a comprehension if you can find a way to refactor your for loop into a simpler format.

While you're still new to comprehensions, I recommend copy-pasting your way from a loop to a comprehension in order to anchor your existing understanding of for loops with your new knowledge of comprehensions.