Python comprehensions

Prettify your code using comprehensions. Beautiful is better than ugly.

·

6 min read

Play this article

For most cases, python comprehensions are a substitute for simple for loops. Are comprehensions faster than for loops? Yes, but that probably won't matter, unless you have absolutely exhausted all the other possible ways to optimize your code. Even then it's a marginal gain.

Instead, use comprehensions to simplify or prettify your code. Beautiful is better than ugly.

There are four types of comprehensions

  • List comprehension
  • Set comprehension
  • Dictionary comprehension
  • Generator comprehension
  • Bonus tip

1. List comprehension

Syntax: [expression with element for element in iterable if condition is met] where the expression and the if clauses are optional. Let's see it in action

# Create a list of alphabets
import string

# A normal for loop would take up 3 lines of code and an indentation
alphabets = []
for alphabet in string.ascii_lowercase:  # string.ascii_lowercase returns >>> 'abcdefghijklmnopqrstuvwxyz'
    alphabets.append(alphabet)

print(alphabets)

"""
OUTPUT:
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
"""

The same code using list comprehension takes up just one line

import string

alphabets = [alphabet for alphabet in string.ascii_lowercase]
print(alphabets)
"""
OUTPUT:
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
"""

Note: The best way to get a list of alphabets is perhaps list(string.ascii_lowercase). The above example is just for illustration purposes. string.ascii_lowercase serves as an example of an iterable

But what if I need only consonants? We can use the optional if clause

vowels = ['a', 'e', 'i', 'o', 'u']                        # Note the if condition below 
alphabets = [alphabet for alphabet in string.ascii_lowercase if alphabet not in vowels]
print(alphabets)
# The if condition filters out the vowels

"""
OUTPUT:
['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't', 'v', 'w', 'x', 'y', 'z']
"""

Now we are only interested in the ordinal value/unicode of the consonants

We can use the ord() function on the alphabet before the for clause as shown below

vowels = ['a', 'e', 'i', 'o', 'u']
unicodes = [ord(alphabet) for alphabet in string.ascii_lowercase if alphabet not in vowels]
print(unicodes)

"""
OUTPUT:
[98, 99, 100, 102, 103, 104, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 118, 119, 120, 121, 122]
"""

Note: We can basically use any function or expressions before the for clause to modify the element at runtime. The most common use case is to write a lambda function in that space


2. Set comprehension

Syntax: {expression with element for element in iterable if condition is met} where the expression and the if clauses are optional. Let's see it in action

# Create a set of unique characters for the below string
string_ = "The quick brown Fox jumps Over the lazy Dog"

# The for loop would look something like this
uniques = set()
for char in string_:
    uniques.add(char)

print(uniques)

"""
OUTPUT:
{'r', 'm', 'F', 'p', 's', 'i', 'u', 'j', 'b', 'c', 'a', 'k', 'o', 'n', 'T', 'x', ' ', 't', 'v', 'z', 'l', 'e', 'w', 'y', 'q', 'h', 'g', 'D', 'O'}
"""

Just like the list comprehension but with curly braces instead of square brackets

uniques = {char for char in string_}
print(uniques)

"""
OUTPUT:
{'r', 'm', 'F', 'p', 's', 'i', 'u', 'j', 'b', 'c', 'a', 'k', 'o', 'n', 'T', 'x', ' ', 't', 'v', 'z', 'l', 'e', 'w', 'y', 'q', 'h', 'g', 'D', 'O'}
"""

Let's say we are not interested in the whitespaces. So we filter it out with the if clause

uniques = {char for char in string_ if char != ' '}
print(uniques)

"""
OUTPUT:
{'r', 'm', 'F', 'p', 's', 'i', 'u', 'j', 'b', 'c', 'a', 'k', 'o', 'n', 'T', 'x', 't', 'v', 'z', 'l', 'e', 'w', 'y', 'q', 'h', 'g', 'D', 'O'}
"""

The set of unique chars still contains a few repetitions in the form of uppercase and lowercase characters. We can convert them to lowercase in the expression clause to avoid the repetition

uniques = {char.lower() for char in string_ if char != ' '}
print(uniques)

"""
OUTPUT:
{'w', 'o', 'b', 'x', 'q', 'v', 'h', 'f', 'j', 'e', 'l', 'r', 'm', 'p', 'c', 't', 'n', 's', 'i', 'y', 'g', 'd', 'a', 'k', 'z', 'u'}
"""

3. Dictionary comprehension

Syntax: {expression for key: expression for valuefor element in iterable if condition is met} where the expression and the if clauses are optional. Let's see it in action

The for loop way

# Create a dictionary of uppercase characters and their unicodes
import string

unicode_map = {}
for char in string.ascii_uppercase:
    unicode_map[char] = ord(char)

print(unicode_map)

"""
OUTPUT:
{
    'A': 65,'B': 66,'C': 67,'D': 68,'E': 69,'F': 70,'G': 71,'H': 72,'I': 73,
    'J': 74,'K': 75,'L': 76,'M': 77,'N': 78,'O': 79,'P': 80,'Q': 81,'R': 82,
    'S': 83,'T': 84,'U': 85,'V': 86,'W': 87,'X': 88,'Y': 89,'Z': 90
}
"""

The dictionary comprehension way

import string

unicode_map = {char: ord(char) for char in string.ascii_uppercase}
print(unicode_map)

"""
OUTPUT:
{
    'A': 65,'B': 66,'C': 67,'D': 68,'E': 69,'F': 70,'G': 71,'H': 72,'I': 73,
    'J': 74,'K': 75,'L': 76,'M': 77,'N': 78,'O': 79,'P': 80,'Q': 81,'R': 82,
    'S': 83,'T': 84,'U': 85,'V': 86,'W': 87,'X': 88,'Y': 89,'Z': 90
}
"""

The expression and the if clauses would work exactly the same as it worked for the list and set comprehensions


4. Generator comprehension

Syntax: (expression with element for element in iterable if condition is met) where the expression and the if clauses are optional. The syntax is very similar to that of a list comprehension just with parenthesis instead of square brackets.

Note: Unlike others, generator comprehensions are a replacement for generator functions rather than for loops.

Let's rewrite the list comprehension example with a generator function

import string

def generate_alphabets():
    for char in string.ascii_lowercase:
        yield char

# The following for loop is to make the generator yield the values.
# It's is NOT the one we need to focus on
for alphabet in generate_alphabets():
    print(f"{alphabet},", end='')

"""
OUTPUT:
a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,
"""

The same function can be written using generator comprehension

import string
generate_alphabets = (char for char in string.ascii_lowercase)

for alphabet in generate_alphabets:
    print(f"{alphabet},", end='')

"""
OUTPUT:
a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,
"""

The expression and the if clause would work exactly the same as the other comprehensions


🥁 Bonus: One line for loops

One more way to make your code concise is to make your for loops into a single line, given that it only has one line in the body as shown below.

Simply write the body right after the :

for char in ['a', 'e', 'i', 'o', 'u']:
    print(char)

# The above for loop can also be written as
for char in ['a', 'e', 'i', 'o', 'u']: print(char)

Did you find this article valuable?

Support Karan Parekh by becoming a sponsor. Any amount is appreciated!