Explain Codes LogoExplain Codes Logo

If/else in a list comprehension

python
list-comprehension
python-best-practices
data-transformation
Anton ShumikhinbyAnton Shumikhin·Oct 3, 2024
TLDR

Use a ternary operator for if/else within a list comprehension:

[new_val if condition else other_val for item in sequence]

For example, tagging numbers as "big" or "small" based on a threshold:

["big" if x > 10 else "small" for x in [8, 12, 10]]

This gives you:

['small', 'big', 'small']

Where new_val is allocated if the condition holds true, else other_val is used across each x in the sequence.

Building conditions in list comprehension

Let's define how to correctly construct conditions within the context of list comprehension to ensure we don't face the Python wrath AKA SyntaxError.

Filtering with list comprehension

Filtering is typically done using a condition towards the end of the list comprehension as shown:

[f(x) for x in sequence if condition]

This removes any x that does not meet the condition. For instance:

[x for x in range(15) if x % 2 == 0] # Even Steven only!

Will return only even numbers between 0 and 14, or as we call them, the "even gang":

[0, 2, 4, 6, 8, 10, 12, 14]

Juggling multiple conditions

You can either nest or combine conditions depending on the complexity of your criteria:

  • Multiple conditions nested:

    [f(x) if condition1 else g(x) if condition2 else default for x in xs] # It's like coding inception!
  • Multiple conditions combined:

    [x for x in sequence if condition1 and condition2] # Best to keep conditions friendly

Syntax: Order in the court!

The for clause must always be the last thing in the comprehension when adding a condition. This helps keep things straight for Python.

Enhance element extraction and transformation

List comprehensions can make your element extraction or data transformation more efficient and concise.

In-list condition checks

You can make your list comprehension more readable by writing the condition before specifying what to iterate over:

[exp if condition else alt for item in iterable]

Here's how you can replace None values in a list with an empty string:

# Replacing None, because 'None' just stands there... doing nothing... [value if value is not None else '' for value in list_maybe_with_none]

Use filter to sanitize lists

filter function can be leveraged to clean up the list before performing list comprehension:

filtered_xs = filter(None, xs) # We only want the ones that matter result = [f(x) for x in filtered_xs]

Python's "truthiness" for the win

In Python, zero (0) is considered False while all other numbers are True. Let's use that in our list comprehension:

# "non-zero" heroes vs "zero" zeroes ['non-zero' if x else 'zero' for x in [-1, 0, 1]]

And the result is:

# Looks like the zeroes are outnumbered! ['non-zero', 'zero', 'non-zero']

Advanced usage and watch-outs

String join() for merging list items

If you ever wanted to merge list items into a single string with different outcomes based on conditions, join() is your friend. Let's create a string of giggles, shall we?

# 'Ha' for the big laughs, 'ha' for the chuckles. "".join(['Ha' if volume > 5 else 'ha' for volume in [3, 6, 4]])

And the laughter ensues:

# Sounds like a fun party! 'hahaHaha'

When in doubt, separate!

When conditions get knotty, separate them into a function for cleaner list comprehension:

def change(x): return f(x) if x else '' results = [change(x) for x in xs]

Lambda: the inline handler

Lambda function can be used in map for inline handling:

# Lambda to the rescue! result = map(lambda x: f(x) if x else '', xs)

Pitfalls to avoid

Complex conditions and deep nesting can be hard to unravel. Keep your comprehensions simple. If things are getting too complex, consider moving out the logic into separate functions or helper methods.