Comprehensions are an extremely powerful part of Python. But what are they and how can they help me?

As you code, you'll notice that certain patterns happen over and over again. Looping through some iterable and create a list is extremely common. Python has a shorthand to make it more straightforward and easier.

A list is a sequence of elements. The elements are placed inside of [ ]. Using a for loop we can create a list by initializing an empty list, then iterating over an object and append elements to the list one by one.

my_numbers = []

for num in range(10):
    my_numbers.append(num)

The range(10) object will provide 0 through 9, so the list in my_numbers is [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]. Inside the for loop we can add in filtering or some processing if needed.

To create a list comprehension, we'll place [] backets with some code inside that will define the sequence of elements.

new_list = [expression 
            for element in iterable       
           ]

The expression is evaluated to create the value to insert into the list. As the for loop is iterating, new values are put through the expression, one by one. We don't need to explicitly code the append() part, Python already understands we're creating a new list.

This list comprehension will create the same list as the for loop above.

my_numbers = [num for num in range(10)]

Changing the value

Let's say we already have a list. [2, 4, 6, 7] and we want a new list with numbers twice as much. We can use num * 2 as the expression and plug that in to build a list comprehension.

my_numbers = [2, 4, 6, 7]
my_numbers = [ num * 2
               for num in my_numbers]

I'm spreading the comprehension into two lines to make it a little easier to see. You can leave it on a single line and that's fine. As your comprehensions get more complicated you might want to spread them out to make them easier to read.

Let's say we have a dictionary of the states where the key is the 2-letter abbreviation and the value is the state name. We can loop over the dictionary by using the items() method. The for loop for abv, name in states.items() is place the key:value pair in abv: name.

Maybe we need to convert the dictionary to a list to pass to a function. What is expected is a list of strings in the format of name (abv) instead.

state_names = [f"{name} ({abv})" 
               for abv, name in states.items()]

Note: See further below for the dictionary and list being used for these examples.

What's nice about this is a list comprehension can be used inline. We can pass it as an argument of a function call. While this does make the line of code a bit longer, we're not storing the list in memory, it's being used right away.

map_setup([f"{name} ({abv})" for abv, name in states.items()])

Filtering the value

Boss comes in and it turns out we shouldn't be mapping all 50 states. Instead, only a set of five states should be added to the list. List comprehensions allow for a conditional portion to be added. We're given a list avail_states that contain the states that should be on the final list.

We'll need to add a condition to our list comprehension that will include only if the name is in a provided list. A simple condition that will check if the name is in the list avail_states and only those that are will be added to the new list.

new_list = [expression 
            for element in iterable
            if condition
           ]

state_names = [f"{name} ({abv})" 
               for abv, name in states.items()
               if name in avail_states]

Two interables

Let's say we have two lists and we want to match each element of the first list with each element of the second list. This is called a Cartesian Product.

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

We'll need to have a loop within a loop.

cartesian_product = []
for a in A:
  for b in B:
    cartesian_product.append( (a, b) )

With a list comprehension, this can make it much shorter.

cartesian_product = [(a, b) for a in A for b in B]

Here's an example on how this be used. With a list of colors and models we can create a new list that contains every color-model match.

colors = ["red", "blue", "black"]
models = ["14", "14 mini", "14 pro"]

inventory = [f"{color} {model}" for color in colors for model in models]

The inventory list end up as this.

[('red', '12'), ('red', '12 mini'), ('red', '12 pro'), ('blue', '12'), ('blue', '12 mini'), ('blue', '12 pro'), ('black', '12'), ('black', '12 mini'), ('black', '12 pro')]`

Examples

Filtering the movies

Using a list of tuples that contain a movie title and year, we can create a new list that contains only the titles of movies that were released prior to 2000.

movies = [(title, year), ...]

pre2k = [title for (title, year) in movies if year < 2000]

Filtering the temperatures

We've been given a set of temperatures and we only want the ones from 30 up to 35.

temperatures = [12, 32, 34, 36, 34, 12, 32]

filtered_temps = [t
                  for t in temperatures
                  if 30 <= t < 35]

Just the numbers please

Starting with a list of strings, we want a new list that contains only integers.

data = ['Mark', 'apple', '7', '10', '23', 'junk']

just_numbers = [int(s)
                for s in data
                if s.isdigit()]

Common elements

Generating a list of students that are in two different lists. There is an option to convert the two lists into sets and then use the .intersection() method of a set. A list comprehension gives us the ability to add more control on what is being added to the new list. If the students were kept in dictionaries a list comprehension gives us control over what is matched. And all this can be done in a single list comprehension.

chem_students = ['Mark', 'Anita', 'Leo', 'Amna']
math_students = ['Amna', 'Zoe', 'Cliff']

in_both = [student
          for student in chem_students
          if student in math_students]

Transposing a matrix

Let's say we want to transpose a matrix, so the rows are columns and columns are rows.

matrix = [
     [1, 2, 3, 4],
     [5, 6, 7, 8],
     [9, 10, 11, 12],
]

new_matrix = [
              [row[i] for row in matrix]
              for i in range(4)]

In this case the expression of the outer list comprehension is itself a list comprehension. The 4 in the range(4) is the length of the inner lists in the matrix.

Even through it's not a list comprehension, I thought this next item is an interesting method of transposing a matrix. The difference is that the inner elements for this will be in tuples and not lists.

new_matrix = list(zip(*matrix))

For either of these transposing I'd recommend putting a comment that the point was to transpose the matrix.

Loading a file

The iterable in the list comprehension doesn't need to be a list. It can be any iterable. We can use the file object. Here we're placing each line of the file as an element in a list. As we're adding it the list, we'll strip off any whitespace from the ends.

with open("raven.txt", 'r') as file:
  poem = [ line.strip() for line in file ]

for line in poem:
    print(line)

Data Source

states = {'AK': 'Alaska', 'AL': 'Alabama', 'AR': 'Arkansas', 'AZ': 'Arizona', 'CA': 'California', 'CO': 'Colorado', 'CT': 'Connecticut', 'DE': 'Delaware', 'FL': 'Florida', 'GA': 'Georgia', 'HI': 'Hawaii', 'IA': 'Iowa', 'ID': 'Idaho', 'IL': 'Illinois', 'IN': 'Indiana', 'KS': 'Kansas', 'KY': 'Kentucky', 'LA': 'Louisiana', 'MA': 'Massachusetts', 'MD': 'Maryland', 'ME': 'Maine', 'MI': 'Michigan', 'MN': 'Minnesota', 'MO': 'Missouri', 'MS': 'Mississippi', 'MT': 'Montana', 'NC': 'North Carolina', 'ND': 'North Dakota', 'NE': 'Nebraska', 'NH': 'New Hampshire', 'NJ': 'New Jersey', 'NM': 'New Mexico', 'NV': 'Nevada', 'NY': 'New York', 'OH': 'Ohio', 'OK': 'Oklahoma', 'OR': 'Oregon', 'PA': 'Pennsylvania', 'RI': 'Rhode Island', 'SC': 'South Carolina', 'SD': 'South Dakota', 'TN': 'Tennessee', 'TX': 'Texas', 'UT': 'Utah', 'VA': 'Virginia', 'VT': 'Vermont', 'WA': 'Washington', 'WI': 'Wisconsin', 'WV': 'West Virginia', 'WY': 'Wyoming'}

avail_states = ['Alaska', 'California', 'Hawaii', 'Oregon', 'Washington']