greendaysbomb
greendaysbomb

Reputation: 404

Effective way to write a series of if statements in Python?

I have a piece of code that I would like to re-write in order for it to be more effective. I am working in Python, and there is not switch-break statement like in Java, but I feel something of a similar structure would work. The piece of code in particular I am talking about is:

def update(arm, array1, array2, array3, array4, array5):
    # If the array is empty, add a zero
    if not array1:
        array1.append(0)
    if not array2:
        array2.append(0)
    if not array3:
        array3.append(0)
    if not array4:
        array4.append(0)
    if not array5:
        array5.append(0)

In general, what is a good way to avoid writing a series of if-statements like this in Python while also going through and checking each of these arrays for this particular function?

Upvotes: 2

Views: 323

Answers (2)

Alain T.
Alain T.

Reputation: 42143

You could define a decorator for any function where you want empty lists to be assigned a zero element. This will allow you to add the initializations without changing anything inside the function.

def noEmptyLists(func):
    def decorated(*args,**kwargs):
        for a in (*args,*kwargs.values()):
            if isinstance(a,list) and not a:
                a.append(0)
        return func(*args,**kwargs)
    return decorated

@noEmptyLists
def update(arm, array1, array2, array3, array4, array5):
    print(arm, array1, array2, array3, array4, array5)


update(1,[],[],[],[],[])

# 1 [0] [0] [0] [0] [0]

Upvotes: 0

DeepSpace
DeepSpace

Reputation: 81594

You can change the function to accept an arbitrary number of arrays with *args, then use a for loop:

def update(arm, *arrays):
    for index, array in enumerate(arrays):
        if not array:
            array.append(0)
        if arm != index:
            array.append(array[-1])

Not only that this is shorter, it works on any arbitrary number of arrays passed to update, so any of update(arm, array1), update(arm, array1, array2), etc will work.

If for some reason you can't change the function's signature, you can still use a loop but you'll have to hardcode the arrays the loop iterates over.

def update(arm, array1, array2, array3, array4, array5,
           array6, array7, array8, array9, array10):
    for index, array in enumerate(array1, array2, array3, array4, array5,
                  array6, array7, array8, array9, array10):
        if not array:
            array.append(0)
        if arm != index:
            array.append(array[-1])

Upvotes: 3

Related Questions