How to achieve two function signatures for `for loop unzip`-CodePudding

Following snippet gives my intention of unzipping t1, t2, t3 or t1, t2 depending upon the task. I know such if-else on for statement doesn't exist but I am wondering is there a workaround for this. Any help or clarification questions are welcome.

def func(task, t1, t2, t3):

    if task == 'abc': # t3=None for this case
        for t1, t2 in zip(t1, t2):
            // do something
    else:
        for t1, t2, t3 in zip(t1, t2, t3):
            // do something

if task == 'abc':
    t3 = None
func(task, t1, t2, t3)

Is there a way we can write a single for loop statement and then unzip the parameters inside of the for loop depending upon the task value. The problem is that when the third parameters is None, it throws the error: TypeError: zip argument #3 must support iteration when task==abc. I want to have a common of do something`.

CodePudding user response：

Use itertools.repeat() to provide a default third list for zip() that produces a continuous stream of a default value.

def func(task, a1, a2, a3 = None):
    if a3 is None:
        a3 = itertools.repeat(None)
    for t1, t2, t3 in zip(a1, a2, a3):
        // do someting

CodePudding user response：

Not clear what the problem is with your existing code, but you could rearrange it like this:

def func(task, *args):
    for items in zip(*args):
        if task == "abc":
            item1, item2 = items
            # do something with them
        else:
            item1, item2, item3 = items
            # do something with these

or this:

def func(task, t1, t2, t3=None):
    if task == "abc":
        lists = [t1, t2]
    else:
        lists = [t1, t2, t3]
    for items in zip(*lists):
    #... as above

or you could use an iterable that endlessly yields None as the default argument:

 from itertools import repeat

 def func(task, t1, t2, t3=repeat(None)):
     for item1, item2, item3 in zip(t1, t2, t3):
         # Now item3 will always be None if t3 is not specified