Set and dictionary comprehensions PREMIUM

In addition to list comprehensions, we also have set and dictionary comprehensions in Python.

Making a set with a for loop

We have a multi-line string (fruits_last_week) in which each line represents the name of a fruit:

fruits_last_week = """
apple
Apple
lime
pear
lime
Lemon
"""

We can split the lines in this string using the string splitlines method:

>>> fruits_last_week.splitlines()
['', 'apple', 'Apple', 'lime', 'pear', 'lime', 'lemon']

Notice that when we split these lines, we get an empty string ('') at the beginning of this list because we have an empty line at the beginning of our multi-line string.

>>> fruits_last_week
'\napple\nApple\nlime\npear\nlime\nLemon\n'

We also have duplicate fruit names in our list.

Let's write a for loop, to build up a new set of unique fruit names (sets in Python don't allow duplicate values):

unique_fruits = set()
for fruit in fruits_last_week.splitlines():
    if fruit:
        unique_fruits.add(fruit.lower())

We're also filtering out that first empty string with a conditional statement (if fruit) which only includes truthy fruits.

>>> unique_fruits
{'lemon', 'apple', 'pear', 'lime'}

Building up a new set with a set comprehension

If you have a for loop that builds up a set and your for loop can be written in the below format, you can turn your for loop into a set comprehension:

new_set = set()
for item in some_iterable:
    if condition_on(item):
        new_set.add(something_with(item))

The way to turn a for loop into a set comprehension is very similar to the way we can turn a for loop into a list comprehension.

To turn this for loop into a set comprehension:

unique_fruits = set()
for fruit in fruits_last_week.splitlines():
    if fruit:
        unique_fruits.add(fruit.lower())

We'll start with the curly braces ({}), because we're building up a set:

>>> unique_fruits = {

The first thing we put after that curly brace is the thing we're adding to new set (fruit.lower() in our case):

>>> unique_fruits = {
...     fruit.lower()

Our looping logic (the for line) comes after that:

>>> unique_fruits = {
...     fruit.lower()
...     for fruit in fruits_last_week.splitlines()

And then if we have a condition, as we do have in this case, that comes last:

>>> unique_fruits = {
...     fruit.lower()
...     for fruit in fruits_last_week.splitlines()
...     if fruit
... }
...

So this set comprehension:

unique_fruits = {
    fruit.lower()
    for fruit in fruits_last_week.splitlines()
    if fruit
}

Is equivalent to this for loop:

unique_fruits = set()
for fruit in fruits_last_week.splitlines():
    if fruit:
        unique_fruits.add(fruit.lower())

A set comprehension looks very similar to a list comprehension except the square brackets ([]) used in a list comprehension are turned into curly braces ({}) in a set comprehension.

A dictionary comprehension is a little different

Writing a dictionary comprehension is a little bit more complex than set comprehension. With a dictionary comprehension, we're not just adding one thing: we have things, a key and a value.

We have a dictionary here (state_to_abbr) that maps Australian state names to state abbreviations:

state_to_abbr = {
    "New South Wales": "NSW",
    "Queensland": "QLD",
    "South Australia": "SA",
    "Tasmania": "TAS",
    "Victoria": "VIC",
    "Western Australia": "WA",
}

And we are building up a new dictionary (abbr_to_state) that does the opposite (it maps abbreviations to states):

abbr_to_state = {}
for state, abbreviation in state_to_abbr.items():
    abbr_to_state[abbreviation] = state

In our new dictionary, we've swapped our key-values pairs (the old values are the new keys).

To make a dictionary comprehension, we'll start with a curly brace ({), just like with our set comprehension before:

>>> abbr_to_state = {

The thing that's different when constructing dictionary comprehension is the first thing we put after that curly brace (the thing we're adding to our new dictionary). We have to add two things, a key and a value. The key (abbreviation) comes first, and the value (state) comes after that, with a colon (:) separating them:

>>> abbr_to_state = {
...     abbreviation: state

Our looping logic (for) comes after that:

>>> abbr_to_state = {
...     abbreviation: state
...     for state, abbreviation in state_to_abbr.items()
... }
...

If we had a condition, that would be the third and last thing in our dictionary comprehension (instead we end our open curly brace).

So this dictionary comprehension:

abbr_to_state = {
    abbreviation: state
    for state, abbreviation in state_to_abbr.items()
}

Does the same thing as this for loop:

abbr_to_state = {}
for state, abbreviation in state_to_abbr.items():
    abbr_to_state[abbreviation] = state

Summary

Just as list comprehensions are a special purpose tool for building up a new list out of an old iterable, set comprehensions and dictionary comprehensions are special purpose tools for building up new sets and dictionaries from old iterables.