Grant
Grant

Reputation: 43

How to combine multiple for loops into a single generator python

I'm on a Python project that uses this for-loop all over the place

for year in yearsindex:
    for month in monthindex:
        for hour in hourindex:
            for node in nodeindex:
                dosomething(year, month, hour, node)

I was wondering if there was a way to combine all of the iterators into a single iterator to be more readable

Something in the form of

for (year, month, hour, node) in combinediterator:
    dosomething(year, month, hour, node)

Upvotes: 1

Views: 168

Answers (2)

user2357112
user2357112

Reputation: 281330

This is itertools.product:

import itertools
for year, month, hour, node in itertools.product(
        yearsindex, monthindex, hourindex, nodeindex):
    dosomething(year, month, hour, node)

You can see that cramming all that onto a single logical line isn't really a readability improvement. There are several ways to make it an improvement. For example, if you can avoid unpacking the tuples the iterator gives you, or if you can put the arguments to itertools.product in a list and unpack them with *args:

for arg_tuple in itertools.product(*indexes):
    dosomething(*arg_tuple)

If the loop body is longer than just one line of dosomething, you also get the benefit of decreased indentation. With a short loop body, that doesn't matter much.

Upvotes: 4

Iron Fist
Iron Fist

Reputation: 10951

Why don't you just wrap it in a function definition in a form of a generator, like so:

>>> l1 = [1,2,3]
>>> l2 = [4,5,6]
>>> l3 = [7,8,9]
>>>
>>> 
>>> def comb_gen(a,b,c):
        for x in a:
            for y in b:
                for z in c:
                    yield (x,y,z)


>>> 
>>> for x,y,z in comb_gen(l1,l2,l3):
        print(x,y,z)


1 4 7
1 4 8
1 4 9
1 5 7
1 5 8
1 5 9
1 6 7
1 6 8
1 6 9
2 4 7
2 4 8
2 4 9
2 5 7
2 5 8
2 5 9
2 6 7
2 6 8
2 6 9
3 4 7
3 4 8
3 4 9
3 5 7
3 5 8
3 5 9
3 6 7
3 6 8
3 6 9

Upvotes: 1

Related Questions