Python: Choose function based on condition in a for loop?-CodePudding

Sorry if the title is a little vague. I'll explain everything in greater detail here. So let's say I have this code:

def function1(k):
    return k * 2


def function2(k):
    return k ** 2


func = 'Square'

for i in range(1, 10):
    if func == 'Multiply':
        function1(i)
    elif func == 'Square':
        function2(i)

How can I modify the code above so that the if statement can go outside the loop? It seems unnecessary to check in every iteration the value of func since it's not going to change inside. the loop. What I'm looking for is something like this:

def function1(k):
    return k * 2


def function2(k):
    return k ** 2


func = 'Square'

if func == 'Multiply':
     f = function1()
elif func == 'Square':
     f = function2()

for i in range(1, 10):
    f(i)

Let me know if something isn't clear enough or if what I'm asking isn't possible.

CodePudding user response：

Store the function in a dict, then you can skip the if statements.

def function1(k):
    return k * 2

def function2(k):
    return k ** 2

d = {"Multiply": function1, "Square": function2}

func = "Square"
f = d[func]
f(3)
# 9

CodePudding user response：

You can do this:

funcs = {'Multiply':function1, 'Square':function2}
theFunc = funcs[func]
for i in range(1, 10):
    x = theFunc(i)
    print(x)

One note in passing: the ^ operator is not taking the square of the argument, but rather performing bitwise xor on it.