Mastering Python Comprehensions: Crafting Efficient and Readable Code

There are a lot of things that Python can do, and one of the things Python is good at is making code more readable and easy to follow. Python comprehension is one of those patterns that will make your code more efficient and readable, unless you go overboard.

So what is comprehension? In technical terms, comprehension is a concise way in Python to create sequences (like lists, sets, and dictionaries) using a single line of code.

So simply put, comprehension will help you create list, set, dict, and generators efficiently and in a more readable format.

Let's start then!

Types of Comprehension

1. List Comprehension

As the name suggests, it is a compact way to create a list. Let's say you want to create a list of all numbers from 1 to 10. How would you do it?

You might say you can just use range(1, 11) to get 1 to 10. But now if you want to get only odd numbers, how would you get it?

To get odd numbers:

odd_nums = [x for x in range(1, 11) if x % 2 != 0]
print(odd_nums)  # Output: [1, 3, 5, 7, 9]

Now with odd numbers, if you want the square of those numbers as well, then you could just do this:

squared_odd_nums = [x**2 for x in range(1, 11) if x % 2 != 0]
print(squared_odd_nums)  # Output: [1, 9, 25, 49, 81]

Now, if you do it the traditional way, this is how your code will look:

squared_odd_nums = []

for num in range(1, 11):
    if num % 2 != 0:
        squared_odd_nums.append(num**2)
print(squared_odd_nums)  # Output: [1, 9, 25, 49, 81]

2. Set Comprehension

No need for more explanation here; set comprehension is the same as list comprehension but instead of composing with [] (square brackets), we will use {} (curly brackets).

squared_odd_nums = {num**2 for num in range(1, 11) if num % 2 != 0}
print(squared_odd_nums)  # Output: {1, 9, 25, 49, 81}

3. Dictionary Comprehension

One more benefit of comprehension is you can easily create a dictionary with one line. Let's say you want a dictionary of all odd numbers and their squared values. How would you do it?

It's easy!

squares = {num: num**2 for num in range(1, 11) if num % 2 != 0}
print(squares)  # Output: {1: 1, 3: 9, 5: 25, 7: 49, 9: 81}

4. Generator Comprehension

Last but not least! Let's see how we can create a generator using comprehension.

Note: Generators look like tuples but they are not.

squared_odd_nums = (num**2 for num in range(1, 11) if num % 2 != 0)
print(type(squared_odd_nums))  # Output: <class 'generator'>
print(list(squared_odd_nums))  # Output: [1, 9, 25, 49, 81]

Now if you really want a tuple, you can do something like this:

squared_odd_nums = tuple(num**2 for num in range(1, 11) if num % 2 != 0)
print(type(squared_odd_nums))  # Output: <class 'tuple'>
print(squared_odd_nums)  # Output: (1, 9, 25, 49, 81)

You might ask how do we use nested loop and conditions with Python comprehension like we do with normal for loop. So here is your answer.

If/Else with Python Comprehension

I showed you earlier how to use if statement with comprehension but what if you want to use if and else both.

Let's see it with similar odd number example that we used earlier. Let' say we want odd number and squared of even number in between 1 to 10. How would we do it?

numbers = [num**2 if num % 2 == 0 else num for num in range(1, 11)]
print(numbers)  # Output: [1, 4, 3, 16, 5, 36, 7, 64, 9, 100]

Here, we have checked that if condition num % 2 == 0 (even number) becomes true then square it otherwise just pass the number (odd number).

Nested loop with Python Comprehension

So now you want to know how the hell we will do the nested loop. So. let's say we want to flatten the 3x3 matrix.

Here is how you would do it with list comprehension.

matrix_3d = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
]
flattened_list = [col for row in matrix_3d for col in row]
print(flattened_list)
# Output: [1, 2, 3, 4, 5, 6, 7, 8, 9]

Here, it will take each row and then iterate through each element of that row.

Now, if you want to use condition with nested loop comprehension here is how you would do it.

matrix_3d = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
]
flattened_list = [col for row in matrix_3d for col in row if col % 2 != 0]
print(flattened_list)
# Output: [1, 3, 5, 7, 9]

You can use if-else same way we use for single loop comprehension. So, that was a guide for comprehension in Python.

Conclusion

From the above, we can conclude that comprehensions are a powerful feature in Python that can help you write cleaner, more efficient code. If you think, "How can this be faster than a normal loop?" check out my blog titled How to read and write sequential data gracefully in Python for a better understanding of that.

Embrace comprehensions to write more Pythonic and maintainable code.