Python N nested loops

Question

So this is an example of what I want to do.

def multi_range(range_func):
    for a in range_func():
        for b in range_func():
            yield a, b

However, what I actually want is something that works for N loops. I would think something like this.

def multi_range(range_func, N, sofar=None, results=None):
    if sofar is None:
        sofar = []

    if results is None:
        results = []

    for a in range_func():
        if N == 1:
            results.append(sofar + [a])
        else:
            multi_range(range_func, N - 1, sofar + [a], results)

    return results

def test_range():
    yield 0
    yield 1

for b in multi_range(test_range, 3):
    print(b)

This correctly outputs the following.

[0, 0, 0]
[0, 0, 1]
[0, 1, 0]
[0, 1, 1]
[1, 0, 0]
[1, 0, 1]
[1, 1, 0]
[1, 1, 1]

The issue here is it has to create the entire list and store that in memory. If for example, test_range were this instead, it would consume a large amount of memory.

def test_range():
    for x in range(10000):
        yield x

(Yes I know it's weird to have a wrapper for range that behaves identically to range. It is just an example)

That is my question. How can I write a function that behaves like this one without storing all results in a list?

Answer 1

Use itertools.product:

>>> from itertools import product
>>> def test_range():
...     yield 0
...     yield 1
...
>>> for b in product(test_range(), repeat=3):
...     print(b)
...
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)

If you're curious how this is implemented, see the sample implementation in the linked doc.