One-Liner For Loops in Python—Code Like a Pro

One-Liner For Loops in Python—Code Like a Pro

You can replace for-loops with one-liner expressions in Python. A one-liner for loop is made possible by utilizing comprehensions.

Python supports four different comprehensions for the main sequence types:

  • List comprehensions
  • Dictionary comprehensions
  • Set comprehensions
  • Generator comprehensions

Today I'm going to show you how each comprehension works. At the end of the day, you know how to improve your code quality by using comprehensions.

List Comprehensions

Let’s start by filtering a list of numbers by leaving out all the negatives. You could use a for loop to solve this task:

numbers = [4, -2, 7, -4, 19]
new_nums = []
for num in numbers:
    if num > 0:
        new_nums.append(num)
print(new_nums)

Output:

[4, 7, 19]

But you can take it to a different level by utilizing list comprehensions:

new_nums = [num for num in numbers if num > 0]
print(new_nums)

Output:

[4, 7, 19]

Cool! You saved four lines of code and slightly improved the readability at the same.

To take home, here is the general structure of the list comprehension:

output_list = [expression for var in input_list if condition]

Notice that the if condition part at the end is optional. For example, let's square the list of numbers:

squares = [n * n for n in numbers]

Dictionary Comprehensions

Python has a similar shorthand for looping through dictionaries as well. It is called the dictionary comprehension. Let's see two examples:

Example 1: Dictionary From a List

Let's create a dictionary from a list of numbers. In the new dictionary, a number is a key and the value is a string representation of that number. In addition, let's leave out all the odd numbers. To solve this task, let's first use a for loop:

nums = [10, 20, 30, 40, 50]
dict = {}

for num in nums:
    if num % 2 == 0:
        dict[num] = str(num)

print(dict)

Output:

{10: '10', 20: '20', 30: '30', 40: '40', 50: '50'}

It works. But let's make it stand out by using dictionary comprehension:

dict = {num: str(num) for num in nums if num % 2 == 0}
print(dict)

Output:

{10: '10', 20: '20', 30: '30', 40: '40', 50: '50'}

That's amazing. Now you know how to turn a list into a dictionary. But that is not what you always want.

Let's see another example where you learn how to transform an existing dictionary into another:

Example 2. Operate on an Existing Dictionary

Let’s use dictionary comprehension to square number values of the dictionary:

data = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}
squared = {k:v*v for (k,v) in data.items()}
print(squared)

Output:

{'a': 1, 'b': 4, 'c': 9, 'd': 16, 'e': 25}

And again, you were able to write a simple version of a for loop with one line of code without forgetting the code readability.

Finally, here is the basic structure of dictionary comprehension to take home:

output_dict = { key:value for (key,value) in dict if condition }

Set Comprehensions

A set comprehension looks similar to list comprehension.

As an example, let's create a set of numbers based on a list of numbers by leaving out all the odd numbers:

numbers = [13, 21, 14, 24, 53, 62] 

filtered_nums = set() 

for num in numbers: 
    if num % 2 == 0: 
        filtered_nums.add(num) 

print(filtered_nums)

Output:

{24, 62, 14}

This is fine. But using comprehensions once again you can save lines of code and make your code stand out:

filtered_nums = {num for num in numbers if num % 2 == 0}
print(filtered_nums)

Output:

{24, 62, 14}

The structure of set comprehensions is similar to that of list and dictionary comprehensions:

output_set = { expression for var in input if condition }

Generator Comprehensions

Finally, let’s go through generator comprehensions. Similar to other comprehensions, generator comprehensions offer a nice shorthand for dealing with generators.

Example: Square all the even numbers in a list and drop all the odd numbers from the list. The for loop approach:

def square_even(numbers):
    for number in numbers:
        if number % 2 == 0:
            yield(number * number)

numbers = [1, 2, 3, 4, 5, 6]
squared_numbers = square_even(numbers)

for number in squared_numbers:
    print(number)

Output:

4
16
36

Once again, this approach works. But especially with generators, there is quite a lot of unnecessary code, such as the separate square_even method. Let's get rid of it with generator comprehensions:

squared_numbers = (num * num for num in numbers if num % 2 == 0)
for number in squared_numbers: 
    print(number)

Output:

4
16
36

And there you go. You just saved 7 lines of code and improved the code quality.

To wrap it up here is the basic structure of a generator comprehension:

output_gen = ( expression for var in input if condition )

Conclusion

In Python, you can turn your for loops into neat little one-liners by using comprehensions.

There are four different comprehensions that are very similar to one another but works for different collection types:

  • List Comprehension
  • Dictionary Comprehension
  • Set Comprehension
  • Generator Comprehension

Thanks for reading. I hope you found this information useful.