Search code examples
pythonlist-comprehensionfor-comprehensionwalrus-operator

How to properly use assignment operator with mod operator in an if statement?


I am trying to use the Walrus operator in python with different if statements and the code that i am trying to replace looks like this:

from time import time

n = 10
numbers = []
results = []

def calculate(x):
    return x ** 2 + 5

t1 = time()
results= [calculate(i) for i in range(n) if (calculate(i)%2) == 0]
t2 = time()
print(f"it took: {t2 - t1} seconds to execute without walrus!")
print(results)

and the expected output should look like this:

it took: 2.5987625122070312e-05 seconds to execute without walrus!
[6, 14, 30, 54, 86]

now if try to replace my code with walrus operator(concept) it does give me either True or 0 in the results if I tried the following:

t1 = time()
results= [result for i in range(n) if (result:= calculate(i) %2  == 0) ]
t2 = time()
print(f"it took: {t2 - t1} seconds to execute with walrus!")
print(results)

Output:

it took: 2.1219253540039062e-05 seconds to execute with walrus!
[True, True, True, True, True]

or:

t1 = time()
results= [result for i in range(n) if ((result:= calculate(i) %2)  == 0) ]
t2 = time()
print(f"it took: {t2 - t1} seconds to execute with walrus!")
print(results)

output:

it took: 2.0742416381835938e-05 seconds to execute with walrus!
[0, 0, 0, 0, 0]

Now I know my logic inside the if statement is wrong and if I did not use comprehension this would probably work as intended but is there any way we can make it work like this?


Solution

  • You're catching the result of the comparison (thus a boolean), not the result of calculate.

    You should do:

    results = [result for i in range(n) if (result:= calculate(i)) %2  == 0 ]