Search code examples
pythonfunctionloopsnestedpython-decorators

Variable number of loops in function (python)


here is an exemple of what I would like to do:

d1 = {'a':1,'b':2,'c':3}
d2 = {'aa':11,'bb':22,'cc':33}
d3 = {'aaa':111,'bbb':222,'ccc':333}

def nLoop(*args):

    n = len(args)

    if n == 1:
        for k0,v0 in args[0].iteritems():
            print k0, v0

    if n == 2:
        for k0,v0 in args[0].iteritems():
            for k1,v1 in args[1].iteritems():
                print k0, v0, k1, v1

    if n == 3:
        for k0,v0 in args[0].iteritems():
            for k1,v1 in args[1].iteritems():
                for k2,v2 in args[2].iteritems():
                    print k0, v0, k1, v1, k2, v2

nLoop(d1,d2,d3)

My question is: are there some ways to do it without the if conditions? Maybe with the use of decorators?


Solution

  • you can pass your variable argument list to itertools.product (with a generator comprenehsion to convert the dicts to their items), then print the flattened results (since product returns tuples of tuples):

    from __future__ import print_function
    import itertools
    
    d1 = {'a':1,'b':2,'c':3}
    d2 = {'aa':11,'bb':22,'cc':33}
    d3 = {'aaa':111,'bbb':222,'ccc':333}
    
    
    def nLoop(*args):
        for t in itertools.product(*(a.items() for a in args)):
            print(*(x for a in t for x in a))
    
    nLoop(d1,d2,d3)
    

    The output of this new nLoop function is identical to yours (if order isn't considered, since dictionary order may change between runs)

    Note that this is a Python 3 compliant solution, but also works with Python 2.