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

October 29, 2022

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`.

>Solution :

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