How To Write Basic List Comprehension In Python

Welcome back, fellow coders!
Today we are going to simplify list comprehension so you can feel confident when writing them.

“List Comprehension” sounds like a massive term when you are first learning how to write the syntax (at least that's how I felt when I started learning) but don’t worry, soon you will be writing them like second nature!

What Is List Comprehension?

Put simply, list comprehension is a powerful tool in Python: It uses a for loop and returns a list object. And ALL this happens in one line. (You can think of list comprehension as the JavaScript equivalent of 'find' 'filter' AND 'map').

Wow... that does sound powerful!
Now let's learn how to write them!

List Comprehension Syntax

When I first started learning list comprehension (after working in JavaScript), I felt like the syntax was backwards and confusing. But if we break the syntax into smaller parts we find that it is quite user-friendly.

In fact, list comprehensions are actually easy to talk through in plain English. For example, the following list comprehension...:

greeting = [print("Hi") for person in my_school] 

... would sound something like this in an English sentence:
"print 'Hi' for each person in my_school"

The above example is a trivial one but it shows the basic layout of our list comprehension syntax. Now let's get more in-depth and discuss each part of the syntax by building our own list comprehension.

Our finished example will look like this:

numbers = [1,2,3,4]

my_numbers = [(num*2) for num in numbers] 

If we were to run this code in our console, we would expect each number to be multiplied by 2, and the following list to be returned when we call my_numbers:

[2, 4, 6, 8]

Let's dive in!

1. Because a list comprehension returns a list, we need to do something with that list (assign it to a variable or return it). In the above example, we assigned it to the variable my_numbers.

my_numbers =  #See the next step

1. Since we are returning a list, we place our list comprehension inside of square brackets.

my_numbers = []

1. The first thing we put inside the square brackets is what we want to return. In our example, we want our num multiplied by 2.

my_numbers = [(num*2)  ]

1. To finish, we write our for loop. In our example, this was for num in numbers where num represents the current item in the list, and numbers is what we are iterating through.

my_numbers = [(num*2) for num in numbers]

And there you have it! That's how simple list comprehensions are to write!
First, assign square brackets to a variable, or place them in a return statement...
Second, add what you want to be returned by the for loop...
And third, finish the list comprehension with the for loop itself.

Tips For Learning List Comprehension

Before we dive deeper into more complex list comprehensions, let's pause to discuss strategies for learning list comprehension.

When I was first learning list comprehension, there were two particular strategies that my professor encouraged me to hone:

1. While learning list comprehension, write out your for loops before writing your list comprehension.

While learning the formatting of list comprehension syntax, I found it extremely helpful to write out my for loops before I wrote my list comprehension.
For example, I might write:

for num in numbers:
    return num * 2

Then I would translate this into my list comprehension:

my_numbers = [(num*2) for num in numbers]

For simple examples like the one above, you might not need to write out your for loop beforehand, but once your code becomes more complicated you will see the value in doing so.

2. The last thing you add to a list comprehension is what you want the for loop to return.

When you are writing out your list comprehension, and are inside the square brackets, give yourself some space and write your for loop before writing what you want it to return. For example, first you would write:

my_numbers = [     for num in numbers]

And then finish off the list comprehension:

my_numbers = [(num*2) for num in numbers]

This will help you return the correct thing when you are working with a more complex list comprehension. It will also allow you to work through the problem in your head in English first. For example, you might say to yourself: "For num in numbers, (now add your return statement) multiply num by 2."

List Comprehensions And Conditional Statements

Hooray! Now you are comfortable with the basic syntax and have a few strategies to practice while learning how to write list comprehensions.
What now?
Before you move on to writing your own code, let's discuss one more important way of writing list comprehension: including conditionals.

Let's continue using our numbers example from above but this time let's make sure that we are only multiplying even numbers. How would we write this in our list comprehension?

By adding our conditional statement after the for loop.
If we were to write out our for loop before writing our list comprehension, our for loop and if statement might look something like this:

# for num in numbers:
#     if num % 2 == 0:
#        num * 2

We would then add our if statement to the list comprehension like so:

numbers = [1,2,3,4]

my_numbers = [(num*2) for num in numbers if num % 2 == 0]

Notice that the "if" comes after the for loop we learned about earlier.
Simply put: add the word if, then add your condition statement.

Example Of List Comprehension In A Project

Now you know the basics and are ready to start writing your own code, but before you do, let's look at a real-world example.

Recently I built a Python project with a CLI menu that stored a list of U.S. States and cities that I want to visit. I can add a State, save it, and then add cities to the State.
I can also list all the States to see what States have already been saved, and then choose one to view details about the cities saved in that State:

(This is my CLI menu. In it, I have chosen to view Hawaii (which I saved prior), and am currently viewing the cities (both are Maui) that are saved to that State of Hawaii. Next to each "city" is the wishlist activity I would like to do in that city/island).

The cities are stored in a database. In order to list all the cities, I need to "fetch" them from the database and make sure they match the State we are currently viewing.
Let's take a look under the hood to see how I accomplished this with list comprehension. In the following example I used list comprehension in city_list to iterate through the cities and find the ones that matched the State:

def list_cities(state):
    cities = City.get_all()
    city_list = [city for city in cities if city.state_id == state.id]
    separator()
    print(f'You are in {state.name}! (Or atleast in your {state.name} Menu...)\n')
    if len(city_list) > 0:
        print("Your cities include:")
        for i, city in enumerate(city_list, start=1):
            print(f'{i}. {city.name}: {city.attraction}')
    else:
        print("Please add a city!")

For explanation purposes, let's take this line by line:
I am currently in a user-facing method called list_cities() that is handed a state object which I will use in a few lines of code to verify that the cities belong to this State:

def list_cities(state):

In cities I am grabbing all my cities in the database using the get_all() method (which is in my City class in a different file, but its main purpose is to get all the cities in my database and turn them into objects that I can use):

    cities = City.get_all()

Now for the list comprehension in city_list where I use a for loop to iterate through cities and check if the foreign key in my city (city.state_id) is equivalent to the primary key in the state object (that was handed to list_cities()). (The state primary key is state.id):

    city_list = [city for city in cities if city.state_id == state.id]

This list comprehension will now return a list of cities that belong to a certain State! In my code I then went on to use that list to format a user-view list of cities in my CLI.

This has been a real-world example of list comprehension. And for those of you who are curious about the rest of my code, let's take a look at it:

    separator()    
    #A line used to help format the user-view.
    print(f'You are in {state.name}! (Or atleast in your {state.name} Menu...)\n')   
    #A print statement used to orient the user while we transition from our State menu to our City menu.
    if len(city_list) > 0:   
    #If our list has at least one city, perform the following actions:
        print("Your cities include:")
        #Since we should never show database id's directly to a user, we use enumerate() to place numbers before each city when we display the cities and their associated attractions.
        for i, city in enumerate(city_list, start=1):    
            print(f'{i}. {city.name}: {city.attraction}')
    else:    
    #If your list does not have any cities in it, print this:
        print("Please add a city!")

Conclusion

Once you become familiar with list comprehension you may find (just like I did) that they are actually kind of fun to write! But beyond how "fun" it is to write them, they are a powerful tool that will take your Python skills to the next level.
So the next time you write a list comprehension remember to break it into parts like we did in this article in order to keep its syntax simple in your head.

Happy coding!

(And a special thanks to my Python instructor Nancy!)