Conditionally swap axes in for loop for multidimensional list Python

Question

I have a group of functions which conditionally switch axes when iterating through a 2-dimensional list. So far I am using an if statement and duplicating the code inside the for loop, and swapping the axes "manually":

def foo(i, axis):

    for j in range(9):
        if axis == 'x':
            doSomething(grid[i][j])
            doSomethingElse(grid[i][j])
            grid[i][j] += 1;
        else:
            doSomething(grid[j][i])
            doSomethingElse(grid[j][i])
            grid[i][j] -= 1;

Is there an "elegant" way to swap the axes which prevents having to re-use code in the for loop for each condition? It would greatly increase readability of the code. Thanks.

Answer 1

you can try this solution

def foo(i, axis):
    for j in range(9):
        value = grid[i][j] if axis == 'x' else grid[j][i]
        doSomething(value)
        doSomethingElse(value)

line 3 is an inline IF statement that works like a simple IF but condensed in only one line

if axis == 'x':
    value = grid[i][j] 
else:
    value = grid[j][i]

and it's better to write only one call to your function and not including that in both blocks of the if because it's more simple to read and easy to maintain