Complex Formatting of Results in List or Dictionary Comprehensions (or Generators)

I love the power that comes with list (or dictionary) comprehensions or generators in Python.

I also despise the code pattern of iterating over object and appending the results to another list.

Let’s keep an example simple with a list comprehension performing a single calculation on each value in the original list:

# Given a list like this:
my_list = [1, 6, 23, 67, 7]

# Good:
results = [x + 1 for x in my_list]

# Bad:
results = []
for x in my_list:
    results.append(x + 1)

In either bits above, results = [2, 7, 24, 68, 8] but you have to admit, that the list comprehension is shorter, easier to write, & faster to understand what’s going on.

Filtering a list comprehension for a single result

If you want to filter out elements in your comprehension, and there’s only one condition you need to apply, you can do something like this:

my_list = [3, 6, 10, 15, 22, 24]

results = [
    x ** 2  # apply this calculation to the value
    for x in my_list
    if x % 3 == 0  # but only when x divides evenly by 3
]

Which returns results = [9, 36, 225, 576] (& tosses 10 & 22 from the results). That can certainly be useful, but what if you want to perform results that have more complex–and multiple–conditions and don’t remove any elements from the list?!

A solution like the one above didn’t cut if for a recent project I was working on. I needed to format the results of a list comprehension based on a set of multiple conditions. (While one approach I started down involved creating multiple list comprehensions–one for each set of output that I needed–and then merging the resulting lists together, the results at best were un-pythonic, & at worst, not only violating the DRY principle but also being terrible to maintain.)

I solved the issue by using a solution I’d done before but never made the connection in my head: Create a function outside the list comprehension & have it perform the necessary conditional formatting to return the results to the list comprehension!

Using functions to handle the complex formatting

def to_my_format(my_number: int) -> int:
    """Perform one of several particular calculations based on the input value given.
    """
    if my_number % 5 == 0:
        return my_number ** 3
    if my_number % 7 == 0:
        return my_number // 2
    return my_number - 1  # If my_number doesn't meet any of the criteria above, do this

my_list = [3, 6, 10, 15, 21, 24]

results = [to_my_format(x) for x in my_list]

Which in case you’re curious, gives us results = [2, 5, 1000, 3375, 3, 23]

So I end up with a list that contains as many elements as the original list (my_list), but one of 3 different calculations are performed on each element depending on the element itself.