How can I write a for-loop that iterates over various forms of iterables?

user_5277458 0
Viewed 52

I'm writing a for-loop that basically iterates over a product of two lists. In most cases, two lists are defined, hence there is no problem with iterating over the product. However, there are also cases where one of the two lists is not defined. Furthermore, there can be a case neither of the lists are defined. Can I change the function bar in the following code snippet, so that it only has a single for-loop that handles various cases without if statements?

import itertools

def foo(a, b):
    print(a, b)

def bar(lista, listb):
    # How can I make this function more concise?
    if lista and listb:
        for a, b in itertools.product(lista, listb):
            foo(a, b)
    elif lista:
        for a in lista:
            foo(a, listb)
    elif listb:
        for b in listb:
            foo(lista, b)
    else:
        foo(lista, listb)

print("Case #1. When both lists are defined.")
lista = [1, 2]
listb = [3, 4]
bar(lista, listb)

print("Case #2. When only lista defined.")
lista = [1, 2]
listb = None
bar(lista, listb)

print("Case #3. When only listb is defined.")
lista = None
listb = [3, 4]
bar(lista, listb)

print("Case #4. When neither of two list are defined.")
lista = None
listb = None
bar(lista, listb)
generic-programming itertools for-loop python-3.x python
2 Answers

Just check ahead of time and replace None values with an acceptable substitute:

import itertools

def bar(lista, listb):
    if lista is None:
        lista = [None]
    if listb is None:
        listb = [None]
    for a, b in itertools.product(lista, listb):
        foo(a, b)

Although IMO, it would be better if the caller had to handle passing valid arguments to bar, and bar assumed that the inputs were always non-None.

Also terminology nitpick, in all cases, lista and listb are defined, they just can be None in some cases. Which is different from being undefined. Since you are talking about parameters to your function, they will always be defined.

user_50144550

Please try below function.

def bar(lista, listb):
    for a, b in itertools.product(*[lista or [None], listb or [None]]):
        foo(a, b)
user_136082280
Related