multirange iterator

Question

I'm doing some boring math with lot of sums over lot of variables. So, instead of nesting for loops I have decided to use something like this:

def combo(*limits):
 if len(limits) == 1:
 for l in range(limits[0]):
 yield [l]
 else:
 for l in range(limits[0]):
 for l2 in combo(*limits[1:]):
 yield [l] + l2
if __name__=="__main__":
 for i, j, k in combo(2, 4, 5):
 print i, j, k

Do you like it? Could it be redone or should it be replaced by something else?

Answer 1

Many interesting loops in Python are better written using itertools. This problem is basically just a kind of product().

from itertools import product
def combo(*limits):
 return product(*map(range, limits))

Answer 2

I suggest to rethink your parameter list. As it stand, if someone tries to call this function without parameters (which is perfectly valid looking at the signature), your function crashes with an IndexError which is not very friendly. Instead, you should declare def combo(term, *remainders): to make it obvious that the function need at least 1 argument. And when call it with 0 arguments, the error will be more explicit:

TypeError: combo() missing 1 required positional argument: 'term'

It also lets you test for the end of the recursion more easily:

def combo(term, *remainders):
 if not remainders:
 for head in range(term):
 yield [head]
 else:
 for head in range(term):
 for tail in combo(*remainders):
 yield [head] + tail

You could also put the test under the first for loop to avoid repeating it twice...

• Eugene: I like that one!
• Eugene: sorry for dis-accepting your answer, but 200_success's answer is better. Still thanks for the parameter hint