Question: Is it possible to call a function inside a list comprehension statement?
Background: List comprehension is a compact way of creating lists. The simple formula is [expression + context].
- Expression: What to do with each list element?
- Context: What elements to select? The context consists of an arbitrary number of
forandifstatements.
For example, the code [x**2 for x in range(3)] creates the list of square numbers [0, 1, 4] with the help of the expression x**2.
Related article: List Comprehension in Python — A Helpful Illustrated Guide
So, can you use a function with or without return value as an expression inside a list comprehension?
Answer: You can use any expression inside the list comprehension, including functions and methods. An expression can be an integer 42, a numerical computation 2+2 (=4), or even a function call np.sum(x) on any iterable x. Any function without return value, returns None per default. That’s why you can even call functions with side effects within a list comprehension statement.
Here’s an example:
[print('hi') for _ in range(10)]
'''
hi
hi
hi
hi
hi
hi
hi
hi
hi
hi
'''You use the throw-away underscore _ because you want to execute the same function ten times. If you want to print the first 10 numbers to the shell, the following code does the trick:
[print(i) for i in range(10)] ''' 0 1 2 3 4 5 6 7 8 9 '''
Let’s have a look at the content of the list you just created:
lst = [print(i) for i in range(10)] print(lst) # [None, None, None, None, None, None, None, None, None, None]
The list contains ten None values because the return value of the print() function is None. The side effect of executing the print function within the list comprehension statement is that the first ten values from 0 to 9 appear on your standard output.
Walrus Operator
Python 3.8 has introduced the walrus operator, also known as the assignment expression. This operator is useful if executing a certain function has side effects that you don’t want. For example, if you have a string creation method inside the list comprehension statement, conditioned by some filtering criterion in the if suffix. Without the walrus operator, Python would execute this same routine multiple times—even though this is highly redundant. You can avoid this redundancy by assigning it to a variable s once using the walrus operator and reusing this exact variable in the expression.
import random
def get_random_string():
return f'sss {random.randrange(0, 100)}'
# Goal: Print all random strings that contain 42
# WRONG
lst = [get_random_string() for _ in range(1000) if '42' in get_random_string()]
print(lst)
# ['sss 74', 'sss 13', 'sss 76', 'sss 13', 'sss 92', 'sss 96', 'sss 27', 'sss 43', 'sss 80']
# CORRECT
lst = [s for _ in range(1000) if '42' in (s := get_random_string())]
print(lst)
# ['sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42', 'sss 42']
With the walrus operator s := get_random_string(), you store the result of the function call in the variable s and retrieve it inside the expression part of the list comprehension. All of this happens inside the list comprehension statement.
I teach these concepts in my exclusive FINXTER email academy—join us, it’s free!