9
\$\begingroup\$

Elements that smell funny: argument unpacking to itertools.product. I'm needing sorted keys (even though I don't care about key order in the final result).

"""
>>> d = {'A': [0, 1, 2],
...      'B': [3, 4],
...      'C': [5, 6, 7]}
>>> for element in dict_product(d):
...     print(element)
{'A': 0, 'B': 3, 'C': 5}
{'A': 0, 'B': 3, 'C': 6}
{'A': 0, 'B': 3, 'C': 7}
{'A': 0, 'B': 4, 'C': 5}
{'A': 0, 'B': 4, 'C': 6}
{'A': 0, 'B': 4, 'C': 7}
{'A': 1, 'B': 3, 'C': 5}
{'A': 1, 'B': 3, 'C': 6}
{'A': 1, 'B': 3, 'C': 7}
{'A': 1, 'B': 4, 'C': 5}
{'A': 1, 'B': 4, 'C': 6}
{'A': 1, 'B': 4, 'C': 7}
{'A': 2, 'B': 3, 'C': 5}
{'A': 2, 'B': 3, 'C': 6}
{'A': 2, 'B': 3, 'C': 7}
{'A': 2, 'B': 4, 'C': 5}
{'A': 2, 'B': 4, 'C': 6}
{'A': 2, 'B': 4, 'C': 7}

"""

from itertools import product

def dict_product(d):
    sorted_keys = sorted(d.keys())
    for element in product(*list(d[k] for k in sorted_keys)):
        yield dict(zip(sorted_keys, element))


if __name__ == "__main__":
    import doctest
    doctest.testmod(verbose=True)
\$\endgroup\$
  • \$\begingroup\$ I had the same need a while back. \$\endgroup\$ – Mathias Ettinger Jan 14 '18 at 17:57
  • \$\begingroup\$ map(dict, product(*[[(k, v) for v in vv] for k, vv in d.items()])) \$\endgroup\$ – Gareth Rees Jan 14 '18 at 19:53
11
\$\begingroup\$

You can use the fact that d.keys() and d.values() are always aligned. This gets rid of both the awkward way of constructing the values as well as the sorting of the keys.

def dict_product(d):
    keys = d.keys()
    for element in product(*d.values()):
        yield dict(zip(keys, element))

The assignment of keys is not strictly necessary, but it avoids the look-up every loop iteration.


Note that list(d[k] for k in sorted_keys) should have been written as a list comprehension directly, instead of calling list on a generator comprehension. So [d[k] for k in sorted_keys].


Your doctest might fail in Python 2 or in Python implementations other than cPython. This is because only recently have dictionary keys become ordered (by insertion time) in (c)Python 3. This is still an implementation detail and not something you should rely upon. Since doctest performs a string comparison, it does not notice that the objects are actually the same.

The only way around this, at least the only one I can see, is to use a more sophisticated testing framework, like unittest. This comes with more code overhead, though.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.