CODEWITHSUNDEEP
pythonif-statementlist-comprehension
Trarbish
Trarbish

Reputation: 371

Pythonic way to use the second condition in list comprehensions

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?

Upvotes: 9

Views: 981

Answers (3)

dheiberg
dheiberg

Reputation: 1914

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
           ]

Upvotes: 8

MSeifert
MSeifert

Reputation: 152587

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.

Upvotes: 8

Uri Goren
Uri Goren

Reputation: 13672

why not use a filter ?

from operator import eq
def myfun(my_list, n, par1=''):
    new_list = ([my_fun2(i,j) for j in range(n)] for i in range(n))
    if par1 != '':
        new_list = filter(eq(par1),new_list)
    return list(new_list)

Upvotes: 1

Related Questions