How to check the "if condition" and assign it on the fly? (Python)

Viewed 102

Say, I want to check if my integer is in the return of the first function or not, if not, then go check the second function, and so on. with these functions

import time

def a_very_expensive_function_1():
    time.sleep(1)
    return [0,1,2]

def a_very_expensive_function_2():
    time.sleep(1)
    return [3,4,5]

# and maybe more functions
# def a_very_expensive_function_3():
#     time.sleep(1)
#     return [6,7,8]
# ...

my_int = 1


if my_int in a_very_expensive_function_1():
    result = a_very_expensive_function_1() + [10] # I have to call`a_very_expensive_function_1()` 2nd times here
elif my_int in a_very_expensive_function_2():
    result = a_very_expensive_function_2() + [10]

You can notice that with this approach, I have to call the first function 2 times, but I don't have to run the 2nd function since the first function is True, (otherwise I have to repeat the 2nd function 2 times just like the first function.)

Now I want to find an alternative, to avoid calling first function 2 times by doing this:

my_int = 1

first_list = a_very_expensive_function_1()
second_list = a_very_expensive_function_2()

if my_int in first_list:
    result = first_list + [10] 
elif my_int in second_list:
    result = second_list + [10]

Now I don't have to call the first function 2 times anymore, but now I have to call a second function (and so on) prior to checking the condition, which I imagine to scale badly if I have more than 2 functions.

Is there an optimal way to tackle these problems?

PS. I am not sure but I feel like walrus operator(:=) can solve this problem? But since it's relatively new feature and my developing environment (Google Colab) is still using Python 3.6 I would prefer more primitive and conventional way to solve this problem.

EDIT: For anyone curious, on Python 3.8+, here is how to use walrus operator to assign on the fly. Again, right now I prefer a more primitive method rather than using a feature that people still not get used to.

import time

def a_very_expensive_function_1():
    time.sleep(1)
    return [0,1,2]

def a_very_expensive_function_2():
    time.sleep(1)
    return [3,4,5]

my_int = 1

if my_int in (result := a_very_expensive_function_1()):
    result += [10] 
elif my_int in (result := a_very_expensive_function_2()):
    result += [10] 

print(result)
# [0, 1, 2, 10]

EDIT2: Add example for more than 2 functions

3 Answers

You can do this:

my_int = 1

first_list = a_very_expensive_function_1()

if my_int in first_list:
    result = first_list + [10] 
else:
    second_list = a_very_expensive_function_2()
    if my_int in second_list:
        result = second_list + [10]

Try it online!

import time

def a_very_expensive_function_1():
    time.sleep(1)
    return [0,1,2]

def a_very_expensive_function_2():
    time.sleep(1)
    return [3,4,5]

def find(n, funcs):
    for e in funcs:
        r = e()
        if n in r:
            return r
    assert False, 'Not found!'

my_int = 1

print(find(my_int, [
    a_very_expensive_function_1,
    a_very_expensive_function_2,
]) + [10])

Output:

[0, 1, 2, 10]

If you're uncomfortable with creating a separate helper function (+ return) then you can do exactly same thing with inline for loop and break:

Try it online!

import time

def a_very_expensive_function_1():
    time.sleep(1)
    return [0,1,2]

def a_very_expensive_function_2():
    time.sleep(1)
    return [3,4,5]


my_int = 1

for e in [
    a_very_expensive_function_1,
    a_very_expensive_function_2,
]:
    r = e()
    if my_int in r:
        break
else:
    assert False, 'Not found!'

print(r + [10])

Output:

[0, 1, 2, 10]

Uhm, can't you just do something like this?

my_int = 1

first_list = a_very_expensive_function_1()   
if my_int in first_list:
    result = first_list + [10] 
else:
    second_list = a_very_expensive_function_2()
    if my_int in second_list:
        result = second_list + [10]

And if you have more than two expensive functions to call:

very_expensive_functions = [
    a_very_expensive_function_1,
    a_very_expensive_function_2,
    a_very_expensive_function_3,
    a_very_expensive_function_4,
    a_very_expensive_function_5
]

my_int = 1

for f in very_expensive_functions:
    f_result = f()
    if my_int in f_result:
        return f_result + [10]
Related