Python Set Comprehensions: How and When to Use Them

Creating Sets With Literals and set()

To create new sets, you can use literals. A set literal is a series of elements enclosed in curly braces. The syntax of a set literal is shown below:

{element_1, element_2,..., element_N}

The elements must be hashable objects. The objects in the literal might be duplicated, but only one instance will be stored in the resulting set. Sets don’t allow duplicate elements. Here’s a quick example of a set:

>>> colors = {"blue", "red", "green", "orange", "green"}
>>> colors
{'red', 'green', 'orange', 'blue'}

>>> colors.add("purple")
>>> colors
{'red', 'green', 'orange', 'purple', 'blue'}

In this example, you create a set containing color names. The elements in your resulting set are unique string objects. You can add new elements using the .add() method. Remember that sets are unordered collections, so the order of elements in the resulting set won’t match the insertion order in most cases.

You can also create a new set using the set() constructor and an iterable of objects:

>>> numbers = [2, 2, 1, 4, 2, 3]

>>> set(numbers)
{1, 2, 3, 4}

In this example, you create a new set using set() with a list of numeric values. Note how the resulting set doesn’t contain duplicate elements. In practice, the set() constructor is a great tool for eliminating duplicate values in iterables.

To create an empty set, you use the set() constructor without arguments:

>>> set()
set()

You can’t create an empty set with a literal because a pair of curly braces {} represents an empty dictionary, not a set. To create an empty set, you must use the set() constructor.

Using for Loops to Populate Sets

Sometimes, you need to start with an empty set and populate it with elements dynamically. To do this, you can use a for loop. For example, say that you want to create a set of unique words from a text. Here’s how to do this with a loop:

>>> unique_words = set()

>>> text = """
... Beautiful is better than ugly
... Explicit is better than implicit
... Simple is better than complex
... Complex is better than complicated
... """.lower()

>>> for word in text.split():
...     unique_words.add(word)
...

>>> unique_words
{
    'beautiful',
    'ugly',
    'better',
    'implicit',
    'complicated',
    'than',
    'explicit',
    'is',
    'complex',
    'simple'
}

In this example, you first create an empty set using set(). You apply lowercase to the text using the .lower() method to make the word extraction case insensitive. Then, you run a loop over a list of words extracted from the text.

As a quick approach, you’ve used the .split() method to split the text into individual words using the space character as a separator.

The loop attempts to add a word to the set in each iteration. However, if a word is already in the set, then any new instance won’t be added. That’s why the resulting set only has one instance of better, for example. The loop is readable and clear, but you can use a set comprehension to make your code even more concise.

Introducing Set Comprehensions

Set comprehensions allow you to build sets with a one-line for loop. If you’re familiar with list comprehensions, then you’ll quickly grasp set comprehensions. Both constructs have similar syntax. The main difference is that set comprehensions use curly braces instead of square brackets.

Here’s the syntax for a set comprehension:

{expression for member in iterable [if condition]}

A set comprehension returns a new set. To build this set, you compute the elements from the items of an input iterable. The syntax includes an optional conditional at the end, which you can use to filter existing collections or generate elements conditionally.

The set comprehension syntax consists of four key elements:

1. Enclosing brackets: Curly braces ({}) are used to define set comprehensions.
2. The comprehension expression: An expression that provides an element in each iteration.
3. The current member: This is the name of the current item or value in the iterable.
4. The iterable: This can be any Python iterable object including a list, tuple, set, generator, or similar types.

The following code shows how you can build a set of unique words using a comprehension:

>>> text = """
... Beautiful is better than ugly
... Explicit is better than implicit
... Simple is better than complex
... Complex is better than complicated
... """.lower()

>>> {word for word in text.split()}
{
    'beautiful',
    'ugly',
    'better',
    'implicit',
    'complicated',
    'than',
    'explicit',
    'is',
    'complex',
    'simple'
}

In this example, you use a set comprehension to extract unique words from the original text. The syntax is concise and clear, improving the readability of your code.

In practice, when you don’t need to transform the data to generate the elements of your set, you can replace the loop from the previous section and the above comprehension with the following code:

>>> set(text.split())
{
    'beautiful',
    'ugly',
    'better',
    'implicit',
    'complicated',
    'than',
    'explicit',
    'is',
    'complex',
    'simple'
}

In this example, you pass the split text directly to the set() constructor and get the same result as with the loop or comprehension. You can do this because the set elements don’t need to run through a transformation.

Comprehensions can also include more than one for clause. When they do, the leftmost for iterates over the outer collection, the next for from left to right iterates over the first nesting level, and so on.

To illustrate, say that you have a list of lists. Each nested list contains numbers, and you want to create a set with their square values. To do this, you can use a comprehension with two for clauses as shown below:

>>> matrix = [
...     [9, 3, 8, 3],
...     [4, 5, 2, 8],
...     [6, 4, 3, 1],
...     [1, 0, 4, 5],
... ]

>>> {value**2 for row in matrix for value in row}
{64, 1, 0, 4, 36, 9, 16, 81, 25}

In this example, the first for loop iterates over the rows of your matrix. The second for loop iterates over the number in each row. As a result, you get a set of square values.

It’s important to note that in this example, instead of having 16 elements, you end up with only nine. This is because set elements are unique. When an element is duplicated, the new instance isn’t added to the final set.

Creating Sets From Iterables

Sometimes, you have an iterable of data and want to create a set using the data items to generate the elements through some data transformation.

For example, say you have a list of tools and libraries. You need to continually check whether a given tool is on the list. To do this, you use membership tests. You also need to make your code case insensitive, so you decide to transform all the words into lowercase letters.

Sets are more efficient than lists in membership tests. Therefore, you also decide to convert your list of tools into a set of lowercase words:

>>> tools = ["Python", "Django", "Flask", "pandas", "NumPy"]

>>> tools_set = {tool.lower() for tool in tools}
>>> tools_set
{'django', 'numpy', 'flask', 'pandas', 'python'}

>>> "python".lower() in tools_set
True
>>> "Pandas".lower() in tools_set
True
>>> "Numpy".lower() in tools_set
True

In the highlighted line, you use a set comprehension while normalizing all values to lowercase to ensure consistent comparisons. Then, you use the in operator to run membership tests and determine whether a given tool is in your original list.

Again, if you don’t need to apply any data transformation to the items of the input iterable, then you don’t have to use a set comprehension. You can use set(iterable) instead. In the above example, the transformation consists of turning the tool names into lowercase.

Transforming Existing Sets

You can also use set comprehensions to quickly transform existing sets. Say that you have a set of user-entered email addresses. You’ve noticed that some addresses include leading and trailing spaces, as well as uppercase letters, and you want to fix this. More important, you’ve noticed that there’s a duplicate email address.

You decide to clean this up with a comprehension:

>>> emails = {
...     " alice@example.org ",
...     "BOB@example.com",
...     "charlie@EXAMPLE.com",
...     "David@example.net",
...     " bob@example.com",
...     "JohnDoe@example.com",
... }

>>> {email.strip().lower() for email in emails}
{
    'alice@example.org',
    'bob@example.com',
    'johndoe@example.com',
    'charlie@example.com',
    'david@example.net'
}

In the highlighted line, you have a set comprehension that removes leading and trailing spaces from each address using the .strip() method. Then, it applies lowercase to all addresses using .lower(). As a result of the cleanup, the duplicate address is eliminated because sets can only contain unique items.

Filtering Items From Sets

Sometimes, you need to filter an existing set and create a new one with elements that meet a given requirement or condition. For example, say that you want to filter your email addresses set to extract only the .com emails. To do this, you can use the following comprehension:

>>> emails_set = {
...     "alice@example.org",
...     "bob@example.com",
...     "johndoe@example.com",
...     "charlie@example.com",
...     "david@example.net",
... }

>>> {email for email in emails_set if email.endswith(".com")}
{'bob@example.com', 'charlie@example.com', 'johndoe@example.com'}

In this example, you use a conditional to create a set that only includes the email addresses on the .com domain.