Short and readable list processing in Python

Viewed 32

In Python, I often work with list processing, which essentially is applying a C-like switch to each element, and doing something with it depending on the case to which it matched.

I have two approaches to how to work with such lists which create 90% of my codebase. I am looking for a possible solution to reduce this clutter and make the code more readable.

The first approach (implementation 1) is based on appending to the list in a loop. While it looks clean and readable it takes too much space for such a simple problem. The second approach (implementation 2) takes less space (probably good enough) but is significantly less readable.

Am I missing some pythonic way of handling lists?

input = [(1, 2), None, (3, 4), (2, 2), None, (7, 1)]

# implementation 1
output = []
for i in input:
    if i is None:
        output.append(i)
    else:
        [a, b] = i
        if (a + b) % 2 == 0:
            output.append((i[0] * 2, i[1]))
        else:
            output.append((i[0], i[1] * 3))

# implementation 2
output = [
    (
        ((i[0] * 2, i[1]) if ((i[0] + i[1]) % 2 == 0) else (i[0], i[1] * 3))
        if i is not None
        else None
    )
    for i in input
]

print(output)
# [(1, 6), None, (3, 12), (4, 2), None, (14, 1)]
3 Answers

For one, we can clean up your implementations a bit by using the sum function.

The first:

output = []
for i in input:
    if i is None:
        output.append(i)
    else:
        if sum(i) % 2 == 0:
            output.append((i[0] * 2, i[1]))
        else:
            output.append((i[0], i[1] * 3))

The second:

output = [
    (
        ((i[0] * 2, i[1]) if sum(i) % 2 == 0) else (i[0], i[1] * 3))
        if i is not None
        else None
    )
    for i in input
]

We can also clean up the first implementation using an elif statement to get something structured more like a switch/case setup.

output = []
for i in input:
    if i is None:
        output.append(i)
    elif sum(i) % 2 == 0:
        output.append((i[0] * 2, i[1]))
    else:
        output.append((i[0], i[1] * 3))

It is also noteworthy that as of version 3.10, Python allows for match/case statements, which is even more analogous to switch/match statements in C.

I recommend using map. A great way to clean up your code is to use pythons built in functional programming tools when possible (map, filter, reduce). In general map can be used when you iterate through a list and perform some operation on each element of the list.

def check(elem):
    if elem is None:
        return elem
    else:
        [a, b] = elem
        if (a + b) % 2 == 0:
            return (a*2, b)
        else:
            return (a, b*3)

inp = [(1, 2), None, (3, 4), (2, 2), None, (7, 1)]
output = list(map(check, inp))

The complexity is imho too high to use implementation 2. My take at it would be to have a function that converts one element, then use list comprehension to turn input to output:

def convert(elem):
   if elem is None:
      return None
   else:
      (a,b) = elem
      return (a*2, b) if (a+b) % 2 == 0 else (a, b*3)

input = [(1, 2), None, (3, 4), (2, 2), None, (7, 1)]
output = [convert(x) for x in input]

print(output)
Related