Introduction

When Python was created, Guido van Rossum and the other developers wanted to make the language as programmer-friendly as they could. Fun fact, if you open up a Python shell and type import this you’ll be met with

Beautiful is better than ugly.
Explicit is better than implicit.
Simple is better than complex.
Complex is better than complicated.
Flat is better than nested.
Sparse is better than dense.
Readability counts.
...

One way in which they followed through on these tenets and allowed programmers to write code “pythonically” is through comprehensions. Comprehensions are clear, concise ways to create lists, dictionaries, sets, and other data structures. Today we’ll be discussing list and dictionary comprehensions.

List Comprehensions

Let's say someone asks you to write a function that takes in a list of ints and outputs a list of the same length, but each number in the list has been squared. Each square in the returned list has to be at the same position as its square root was in the parameter list. One way to solve this problem is as follows:

Without a List Comprehension

def make_list_of_squares(ints):
    """
    This function takes in a list of ints and returns a list of their squares, in the same order.
    """
    squares = []

    for int in ints:
        square = int ** 2
        squares.append(square)

    return square

With a List Comprehension
Using a list comprehension, we can write the same function in one line, as follows:

def make_list_of_squares(ints):
    """
    This function takes in a list of ints and returns a list of their squares, in the same order.
    """
    return [int ** 2 for int in ints]

Let’s break down that line, as its pretty dense. A list comprehension is broken down into three parts: an expression, a for loop, and an optional conditional statement. Using these three facets, we can create lists in just one line! Here’s the anatomy of a list comprehension:

[..expression....for loop....(optional) if statement..]

The Optional If Statement
In the example we gave above, of making the list of square roots with a list comprehension, we used 2 of the 3 components of the list comprehension, omitting the optional if statement:

• expression: int ** 2
• for loop: for int in ints

If we wanted to make a list of all squares of even numbers in the input list, we could do the following:

[ int ** 2 for int in ints if int % 2 == 0 ]

Here, we have all three components

• expression: int ** 2
• for loop: `for int in ints
• (optional) if statement: if int % 2 == 0

More Examples of List Comprehensions

List comprehensions can get even more powerful when applying other functions in your programs. For example, if I had a list of User objects (User.objects.all()) and wanted to make a list of the user’s name, only if they’re an admin user, I could do:

[ user.get_name() for user in User.objects.all() is user.is_admin() ]

Concluding List Comprehensions

List comprehensions are a very powerful tool that Python comes with. They can make complex calculations and list constructions simple and more readable. I hope this becomes another tool in your arsenal to make your code more functional and more readable!

Dictionary Comprehensions

Python also allows you to make dictionary comprehensions. They’re similar in construction to list comprehensions, here’s the anatomy of a dictionary comprehension

[..key : value....for (key, value) in collection....(optional) if statement based on (key, value)..]

Reversing a Dictionary Using a Comprehension
A classic example of using a dictionary comprehension is to reverse a dictionary. For this example, note that dict.items() returns a list of (key, value) tuples that comprise the dictionary.

def reverse_dictionary(my_dict):
    """
    This function takes in a dictionary and returns the reversed version of it, such that 
    the values are mapped to the keys.
    """
    return { value : key for (key, value) in my_dict.items() }

Conclusion on Dictionary Comprehensions and List Comprehensions

Now, let’s try putting everything together by writing a function that takes in a list of string and returns a dictionary that has the keys as letters and values as lists of strings in the input list that start with that letter. This would look something like this without a dictionary comprehension:

def make_first_letter_dictionary(lst):
    """
    This function takes in a list of strings and returns a map that associates 
    letters with a list of strings in the input list that start with that letter. For 
    example, if our input list was 

    [ 'apple', 'ball', 'avocado', 'bagel' ] 

    our returned dictionary would be 

    { 'a' : ['apple, 'avocado'], 'b' : ['ball', 'bagel']}
    """

    result = {} 

    for str in lst:
        key = str[0]
        # If key not in dictionary, initialize the list it maps to
        if key not in result:
            result[key] = []
        result[key].append(str)

    return result

Using both list comprehensions and dictionary comprehensions, we can write the function in one line:

def make_first_letter_dictionary(lst):
    """
    Using comprehensions to solve the problem above.
    """
    return { str[0] : [ s for s in lst if s[0] == str[0] ] for str in lst}

Let’s break down what’s going on here. For the dictionary comprehension, we’re looping through the strings in lst and mapping the first letter of each string to a list comprehension. In that list comprehension, we’re looping through each string in lst and extracting all strings that start with the current letter we’re on.

Choosing When to Use Comprehensions
Now, it’s worth noting that even though this version of the function is stylistically pleasing, it is not as efficient as the first implementation. The comprehension version if looping through the input list once to get the keys and again to make each inner list. As a result, this version has a runtime complexity of O(n^2), while the first implementation makes one pass through the list, resulting in a O(n) runtime, where n = len(lst). As such, I would advise choosing a comprehension when it makes both stylistic and functional sense. However clean it may be, it’s not always the most efficient choice.