How can I make a for-loop pyramid more concise in Python? [duplicate]

Question

In solid mechanics, I often use Python and write code that looks like the following:

for i in range(3):     for j in range(3):         for k in range(3):             for l in range(3):                 # do stuff 

I do this really often that I start to wonder whether there is a more concise way to do this. The drawback of the current code is: if I comply with PEP8, then I cannot exceed the 79-character-limit per line, and there is not too much space left, especially if this is again in a function of a class.

Answer 1

Based on what you want to do, you can use the itertools module to minimize the for loops (or zip).In this case itertools.product would create what you have done with the 4 loops:

>>> list(product(range(3),repeat=4)) [(0, 0, 0, 0), (0, 0, 0, 1), (0, 0, 0, 2), (0, 0, 1, 0), (0, 0, 1, 1),  (0, 0, 1, 2), (0, 0, 2, 0), (0, 0, 2, 1), (0, 0, 2, 2), (0, 1, 0, 0),  (0, 1, 0, 1), (0, 1, 0, 2), (0, 1, 1, 0), (0, 1, 1, 1), (0, 1, 1, 2),  (0, 1, 2, 0), (0, 1, 2, 1), (0, 1, 2, 2), (0, 2, 0, 0), (0, 2, 0, 1),  (0, 2, 0, 2), (0, 2, 1, 0), (0, 2, 1, 1), (0, 2, 1, 2), (0, 2, 2, 0),  (0, 2, 2, 1), (0, 2, 2, 2), (1, 0, 0, 0), (1, 0, 0, 1), (1, 0, 0, 2),  (1, 0, 1, 0), (1, 0, 1, 1), (1, 0, 1, 2), (1, 0, 2, 0), (1, 0, 2, 1),  (1, 0, 2, 2), (1, 1, 0, 0), (1, 1, 0, 1), (1, 1, 0, 2), (1, 1, 1, 0),  (1, 1, 1, 1), (1, 1, 1, 2), (1, 1, 2, 0), (1, 1, 2, 1), (1, 1, 2, 2),  (1, 2, 0, 0), (1, 2, 0, 1), (1, 2, 0, 2), (1, 2, 1, 0), (1, 2, 1, 1),  (1, 2, 1, 2), (1, 2, 2, 0), (1, 2, 2, 1), (1, 2, 2, 2), (2, 0, 0, 0),  (2, 0, 0, 1), (2, 0, 0, 2), (2, 0, 1, 0), (2, 0, 1, 1), (2, 0, 1, 2),  (2, 0, 2, 0), (2, 0, 2, 1), (2, 0, 2, 2), (2, 1, 0, 0), (2, 1, 0, 1),  (2, 1, 0, 2), (2, 1, 1, 0), (2, 1, 1, 1), (2, 1, 1, 2), (2, 1, 2, 0),  (2, 1, 2, 1), (2, 1, 2, 2), (2, 2, 0, 0), (2, 2, 0, 1), (2, 2, 0, 2),  (2, 2, 1, 0), (2, 2, 1, 1), (2, 2, 1, 2), (2, 2, 2, 0), (2, 2, 2, 1),  (2, 2, 2, 2)] 

And in your code you can do :

for i,j,k,l in product(range(3),repeat=4):     #do stuff 

This function is equivalent to the following code, except that the actual implementation does not build up intermediate results in memory:

def product(*args, **kwds):     # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy     # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111     pools = map(tuple, args) * kwds.get('repeat', 1)     result = [[]]     for pool in pools:         result = [x+[y] for x in result for y in pool]     for prod in result:         yield tuple(prod) 

Edit:As @ PeterE says in comment product() can be used even if the ranges have differing length :

product(range(3),range(4),['a','b','c'] ,some_other_iterable) 

Answer 2

The idea to use itertools.product is a good one. Here's a more general approach that will support ranges of varying sizes.

from itertools import product  def product_of_ranges(*ns):     for t in product(*map(range, ns)):         yield t  for i, j, k in product_of_ranges(4, 2, 3):     # do stuff