List comprehensions

List comprehensions provide a concise way to create lists in Python, combining a for loop and list creation into a single line. Python also offers set and dictionary comprehensions, which we will meet later.

Syntax

[expression for item in iterable]

• expression - the operation applied to each item.
• item - the loop variable, holding one element at a time.
• iterable - the object we are iterating over (e.g. a list, tuple, or dictionary).

In plain terms: "for every item in the iterable, evaluate the expression, and collect the results into a new list". We could achieve the same with a for loop - a list comprehension is simply more concise.

Consider a list of numbers where we want to square every element and store the results in a new list. Using a for loop:

numbers = [2, 4, 3, 5]
squared_numbers = []  # create new list to hold output

for number in numbers:
    squared_numbers.append(number**2)

print(squared_numbers)  # Output: [4, 16, 9, 25]

Using a list comprehension, we achieve the same result in a single line:

numbers = [2, 4, 3, 5]
squared_numbers = [x**2 for x in numbers]
print(squared_numbers)  # Output: [4, 16, 9, 25]

Breaking this down:

• x**2 is the expression - it squares the element on each iteration.
• x is the loop variable. We could name it anything we like.
• for x in numbers defines the iteration over the numbers list.
• The square brackets ([]) mean the results are collected into a new list, which we assign to squared_numbers.

Here are some more examples of list comprehensions with different expressions and iterable types.

Example 1 - String indexing

We can extract the first letter of each name to create a list of initials:

names = ['Alice', 'Bob', 'Charlie', 'Diana']
initials = [name[0] for name in names]
print(initials)  # Output: ['A', 'B', 'C', 'D']

If we are unsure whether the first letter is capitalised, we can chain a method call in the expression:

names = ['alice', 'Bob', 'charlie', 'Diana']
initials = [name[0].upper() for name in names]
print(initials)  # Output: ['A', 'B', 'C', 'D']

Here .upper() is called on the character returned by name[0], ensuring every initial is upper-case.

Example 2 - String methods

Given a list of email addresses, we can extract the username (the text before @) for each:

emails = ['[email protected]', '[email protected]', '[email protected]']
usernames = [email.split('@')[0] for email in emails]
print(usernames)  # Output: ['john.doe', 'jane.smith', 'info']

The expression email.split('@')[0] splits each email string at @ into a list of two parts, then takes the first element (the username).

Example 3 - Dictionary key-value pairs

We can iterate over both keys and values of a dictionary using .items():

numbers_dict = {0: 1, 1: 2, 2: 3}
transformed = [k*v for k, v in numbers_dict.items()]
print(transformed)  # Output: [0, 2, 6]

Here k, v unpacks each key-value pair from numbers_dict.items(), and the expression k * v multiplies the key by its corresponding value.