Python List Comprehensions Are More Powerful Than You Might Think
Write better list comprehensions with the help of these unknown features and tricks
Python’s list comprehensions (and generators) are an awesome feature that can greatly simplify your code. Most of the time, however, we only use them to write a single for loop, maybe with the addition of one if conditional, and that's it. If you start poking around a bit though, you will find out that there are many more features of Python's comprehensions that you don't know about, but can learn a lot from…
Multiple Conditionals
We know that we can use if conditional to filter results of a list comprehension and with simple comprehensions, single if is usually sufficient. What if you wanted a nested conditional, though?
It’s possible to build a nested conditional using “conditional expressions,” or as it’s generally called a ternary operator. It’s not exactly a pretty solution, so you will have to decide whether the few saved lines are worth the nasty one-liner.
Apart from using complex conditionals, it’s also possible to stack multiple ifs in a comprehension. Here’s the code:
Looking at the expanded code above, it doesn’t really make much sense to write it this way, but the syntax allows it.
One reason why you might want to use it is for readability purposes, as shown below:
Avoid Repeated Evaluation
Let’s say you have comprehension that calls an expensive function both in the conditional and in the loop body. It may look like the following:
This is inefficient as it doubles the computation time, but what can we do about it? Nested comprehensions to the rescue!
I want to highlight that the above is not a double loop. In this example, we build a generator inside list comprehension which is consumed by the outer loop. If you find this hard to read, then the alternative would be to use the walrus operator. Here’s the code:
Here func is called only once, creating a local variable y which can be used in other parts of the expression.
Handling Exceptions
Even though list comprehensions are usually used for simple tasks — such as calling a function on each element of the list — there are situations where exceptions might be thrown inside the comprehension. However, there’s no native way of handling an exception inside a list comprehension, so what can we do about it?
We need a handler function to catch an exception inside a comprehension. Here we create a function catch that takes a function and its arguments. If an exception is thrown inside catch, then the exception is returned.
This is not an ideal solution, considering that we need a helper function, but it’s the best we can do as the proposal ( PEP 463), which tried to introduce a syntax for this, got rejected.
Breaking the Loop
Another limitation of list comprehensions is the inability to break the loop. While not possible natively, we can implement a little hack that solves the problem:
The first example above uses a little know feature/behavior of iter function. The iter(callable, sentinel) returns an iterator that "breaks" iteration once callable function value equals a sentinel value. When the inner iter returns the sentinel ( 4 in the example), the loop automatically stops.
This is not very readable, so you can instead take advantage of the great itertools module and the takewhile function, as shown in the second example.
As a side note — if you thought breaking a loop in list comprehensions was possible, then you’d be correct. Until Python 3.5, you could use a helper function to raise StopIteration inside list comprehensions, that was, however, changed with PEP 479.
Tricks (and Hacks)
In the previous sections, we’ve seen some obscure features of list comprehensions that might or might not be very useful in day-to-day coding. So, let’s now look at some tricks (and little hacks) that you can use right away.
While plain, vanilla list comprehensions are very powerful, they become even better when paired with libraries such itertools (see the previous section) or its extension more-itertools.
Let’s say you need to find runs of consecutive numbers, dates, letters, booleans, or any other orderable objects. You can solve this elegantly by pairing consecutive_groups from more-itertools with a list comprehension. Here’s the code:
Here we have a list of dates, some of which are consecutive. We pass the dates to the consecutive_groups function using ordinal values of the dates for ordering. We then collect returned groups into a list using comprehension.
Computing accumulating sums of numbers is very easy in Python — you can just pass a list to itertools.accumulate, and you get back the sums. What if we wanted to undo the accumulation, though?
With help of more_itertools.pairwise it's pretty simple!
As was mentioned earlier, the new-ish walrus operator can be used with list comprehensions to create a local variable. That can be useful in many situations. One such situation is with any() and all() functions:
Python’s any() and all() functions can verify whether any or all values in some iterable satisfy certain conditions. What if you, however, want to also capture the value that caused any() to return True (so-called "witness") or the value that caused all() to fail (so-called "counterexample")?
Both any() and all() use short-circuiting to evaluate the given expression. This means they stop the evaluation as soon as they find the first "witness" or "counterexample" respectively. Therefore, with this trick, the variable created by the walrus operator will always give us the first "witness"/"counterexample."
Closing Thoughts
Many features and tricks demonstrate the possibilities and limits of list comprehensions. Learning these intricacies is — in my opinion — a good way of gaining a better understanding of particular language features, even if it’s not really useful in daily coding. On top of that, it’s fun.
With that said, I hope you learn something here and be aware that if you decide to use complex conditionals or loop breaks in your list comprehensions, your coworkers might end up hating you.
Want to Connect?This article was originally posted at martinheinz.dev.
