python: how do I pull the If else statement into this loop when lambda is involved?

Question

I have many functions that sometimes should loop sometimes shouldn't. I want to build in an option to tell it to loop. Here is my code. You can see that every function has the if... else statement. Is there a way to put the if... else statement into the method loop? So I don't have to repeat those lines for every function?

import inspect
def test1(b,shouldloop=False):
 if not shouldloop:
 a = b + 1
 print a
 else:
 loop(lambda z: test1(z) , inspect.currentframe().f_code.co_name)
def test2(d,e,shouldloop=False):
 if not shouldloop:
 a = d * e
 print a
 else:
 loop(lambda z: test2(z,e) , inspect.currentframe().f_code.co_name)
def test3(g,h,i,shouldloop=False):
 if not shouldloop:
 a = g **2 - h + i
 print a
 else:
 loop(lambda z: test3(z,h,i) , inspect.currentframe().f_code.co_name)
def loop(function,c):
 x = [1,2,3,4]
 for i in x:
 function (i)
 print c

so sometimes I just need the function test1 to not loop like so...

test1(2)

output:

3

and sometimes I need it to loop like so...

test1('',True)

output:

2
test1
3
test1
4
test1
5
test1

The other test functions have more variables. But I will be using them the same way.

• Sayse: I don't understand your logic flow here, what is this code supposed to output? Also, instead of using 'yes' for your shouldloop you should just use a boolean value
• FMarazzi: I agree with @Sayse, that is quite complicated. Can you provide a more clear explanation of your code?
• r.ook: Since you're doing 6 separate things in the 3 sets of if... else... blocks, it's probably best to leave it like this. I fear any other attemps to pass a set of if... else... block that does different things into different functions is just gonna be a hot mess.

Answer 1

You can rewrite loop as a decorator to add it around arbitrary functions:

import functools
def loopable(fnc): # decorator, takes any function
 @functools.wraps(fnc)
 def fncloop(*args, shouldloop=False, **kwargs): # take loop switch and arguments for the actual function
 if shouldloop: # loop logic
 x = [1,2,3,4]
 for i in x:
 fnc(i, *args[1:], **kwargs) # replace first argument with i
 print(fnc.__name__)
 else:
 fnc(*args, **kwargs) # call function unconditionally
 return fncloop # return wrapped function

Apply it when defining a function, which just contains the loop body:

@loopable
def test1(b):
 a = b + 1
 print(a)

Any decorated function now supports shouldloop:

>>> test1(2)
3
>>> test1(2, shouldloop=True)
2
test1
3
test1
4
test1
5
test1

---

For Python 2, you have to fetch the argument from kwargs explicitly:

import functools
def loopable(fnc):
 @functools.wraps(fnc)
 def fncloop(*args, **kwargs):
 if kwargs.pop('shouldloop', False):
 x = [1,2,3,4]
 for i in x:
 fnc(i, *args[1:], **kwargs) # replace first argument with i
 print(fnc.__name__)
 else:
 fnc(*args, **kwargs)
 return fncloop