Search code examples
pythoniterationenumerationpython-itertools

Conditionally enumerate through all elements of a dict in python


I have a dict as such:

config = {
          'CONF_A': ['a1', 'a2', 'a3'],
          'CONF_B': ['b1', 'b2'],
          'CONF_C': {
                     'CONF_C1': ['c1'],
                     'CONF_C2': ['c21','c22']
                    }
          }

I need to enumerate through all values of CONF_A and CONF_B, firstly using all the values in CONF_C1, and then all the values in CONF_C2. e.g:

('a1', 'b1', 'c1')
('a1', 'b2', 'c1')
('a2', 'b1', 'c1')
('a2', 'b2', 'c1')
('a3', 'b1', 'c1')
('a3', 'b2', 'c1')
.... same for 'c21'
.... same for 'c22'

I also need to have the key for each of these too. I can get this for a config consisting of just CONF_A and CONF_B (see below), but I'm not sure how to also enumerate each of the CONF_C elements.

>>> for x in itertools.product(*conf.itervalues()):
        for k, v in itertools.izip(conf,x):
            print k,v
...        
    CONF_A a1
    CONF_B b1
    CONF_A a1
    CONF_B b2
    CONF_A a2
    CONF_B b1
    CONF_A a2
    CONF_B b2
    CONF_A a3
    CONF_B b1
    CONF_A a3
    CONF_B b2

N.B. The keys for CONF_C elements should be the 'CONF_C1' and 'CONF_C2'

---- edit ----

@Patrick.Haugh provided this answer for the first part of the question:

 for x in product(config['CONF_A'], config['CONF_B'], chain.from_iterable(config['CONF_C'].values())):

This works for getting the cartesian product, however I still need to get the keys back in. When I do the following:

for k, v in itertools.izip(config, x):

I get the keys CONF_A, CONF_B, and CONF_C. However the values are in reverse order, e.g.

{'CONF_A': c1, 'CONF_B': b1, 'CONF_C': 'a1'}

Additionally, I'm trying to figure out how to ensure that the keys for C match up with the CONF_C values, i.e.

{'CONF_A': a1, 'CONF_B': b1, 'CONF_C1': 'c1'}
{'CONF_A': a1, 'CONF_B': b1, 'CONF_C2': 'c21'}

Solution

  • Use itertools.product with itertools.chain

    from itertools import product, chain
    
    list(product(config['CONF_A'], config['CONF_B'], chain.from_iterable(config['CONF_C'].values())))
    

    gives us

    [('a1', 'b1', 'c1'), ('a1', 'b1', 'c21'), ('a1', 'b1', 'c22'), ('a1', 'b2', 'c1'), 
     ('a1', 'b2', 'c21'), ('a1', 'b2', 'c22'), ('a2', 'b1', 'c1'), ('a2', 'b1', 'c21'), 
     ('a2', 'b1', 'c22'), ('a2', 'b2', 'c1'), ('a2', 'b2', 'c21'), ('a2', 'b2', 'c22'), 
     ('a3', 'b1', 'c1'), ('a3', 'b1', 'c21'), ('a3', 'b1', 'c22'), ('a3', 'b2', 'c1'), 
     ('a3', 'b2', 'c21'), ('a3', 'b2', 'c22')]
    

    Edit:

    a_gen = (('CONF_A', a) for a in config['CONF_A'])
    b_gen = (('CONF_B', b) for b in config['CONF_B'])
    c_gen = ((k, e) for k, v in config['CONF_C'].items() for e in v)
    for i in product(a_gen, b_gen, c_gen):
        print(i)
    

    gives us

    (('CONF_A', 'a1'), ('CONF_B', 'b1'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a1'), ('CONF_B', 'b1'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a1'), ('CONF_B', 'b1'), ('CONF_C2', 'c22'))
    (('CONF_A', 'a1'), ('CONF_B', 'b2'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a1'), ('CONF_B', 'b2'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a1'), ('CONF_B', 'b2'), ('CONF_C2', 'c22'))
    (('CONF_A', 'a2'), ('CONF_B', 'b1'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a2'), ('CONF_B', 'b1'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a2'), ('CONF_B', 'b1'), ('CONF_C2', 'c22'))
    (('CONF_A', 'a2'), ('CONF_B', 'b2'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a2'), ('CONF_B', 'b2'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a2'), ('CONF_B', 'b2'), ('CONF_C2', 'c22'))
    (('CONF_A', 'a3'), ('CONF_B', 'b1'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a3'), ('CONF_B', 'b1'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a3'), ('CONF_B', 'b1'), ('CONF_C2', 'c22'))
    (('CONF_A', 'a3'), ('CONF_B', 'b2'), ('CONF_C1', 'c1'))
    (('CONF_A', 'a3'), ('CONF_B', 'b2'), ('CONF_C2', 'c21'))
    (('CONF_A', 'a3'), ('CONF_B', 'b2'), ('CONF_C2', 'c22'))