Pythonic way to use the second condition in list comprehensions

Question

Let's assume the following function:

def myfun(my_list, n, par1=''):
    if par1 == '':
        new_list = [[my_fun2(i,j) for j in range(n)] for i in range(n)]
    else:
        new_list = [[my_fun2(i,j) for j in range(n)] for i in range(n) if my_fun2(i,n) == par1]
    return new_list

As you can see, there are two different scenarios depending on par1. I do not like that line 3 and line 5 are almost identical and do not follow the DRY (Don't Repeat Yourself) principle. How can this code be improved?

Answer 1

You could choose the condition function dynamically by using a function that just returns True in the first case and one that actually compares the my_fun2 result with par1 in the second case:

def myfun(my_list, n, par1=''):
    if par1 == '':
        cond = lambda x, y: True
    else:
        cond = lambda i, n: my_fun2(i, n) == par1
    return [[my_fun2(i,j) for j in range(n)] for i in range(n) if cond(i,n)]

Or by replacing the outer loop with a generator expression in case par1 isn't an empty string:

def myfun(my_list, n, par1=''):
    if par1 == '':
        outer = range(n)
    else:
        # a conditional generator expression
        outer = (i for i in range(n) if my_fun2(i,n) == par1)
    return [[my_fun2(i,j) for j in range(n)] for i in outer]

However don't let DRY make the function harder to read, maintain or debug. I, personally, think that your approach is fine (and probably faster) and you probably shouldn't change anything.

Answer 2

This might work:

new_list = [[my_fun2(i,j) for j in range(n)] for i in range(n) if par1 == '' or my_fun2(i,n) == par1]

So used like this:

def myfun(my_list, n, par1=''):
    return [
               [my_fun2(i,j) for j in range(n)]
               for i in range(n) if par1 == '' or my_fun2(i,n) == par1
           ]