Search code examples
pythonpython-3.xdictionarykeytraversal

How to navigate a dict by list of keys?


Starting with a nested dictionary like this:

my_dict = {"North America" : { "USA" : { "Virginia" : ["Norfolk","Richmond","Charlottesville"], "New York": ["Albany"]}, "Canada" : {"Saskatchewan": ["Saskatoon"], "New Brunswick":["Moncton","Saint John"]}}}

print(my_dict)
North America
    USA
        Virginia
            ['Norfolk', 'Richmond', 'Charlottesville']
        New York
            ['Albany']
    Canada
        Saskatchewan
            ['Saskatoon']
        New Brunswick
            ['Moncton', 'Saint John']

How can I take a list of strings like [key1, key2] and programmatically return the nested object my_dict[key1][key2] when I don't know how many keys there will be? Examples:

keys = ['North America', 'USA']
print(my_dict.???)
    Virginia
        ['Norfolk', 'Richmond', 'Charlottesville']
    New York
        ['Albany']

keys = ['North America', 'Canada', 'Saskatchewan']
print(my_dict.???)
    ['Saskatoon']

keys = ['North America', 'Canada']
print(my_dict.???)
    Saskatchewan
        ['Saskatoon']
    New Brunswick
        ['Moncton', 'Saint John']

Assume the "path" to the keys exists, but don't assume anything about depth in the structure.


Solution

  • You can throw functools.reduce onto it.

    >>> from functools import reduce
    >>>
    >>> my_dict = {"North America" : { "USA" : { "Virginia" : ["Norfolk","Richmond","Charlottesville"], "New York": ["Albany"]}, "Canada" : {"Saskatchewan": ["Saskatoon"], "New Brunswick":["Moncton","Saint John"]}}}
    >>> keys = ['North America', 'Canada']
    >>> reduce(dict.get, keys, my_dict)
    >>> {'New Brunswick': ['Moncton', 'Saint John'], 'Saskatchewan': ['Saskatoon']}
    

    ... or write your own function.

    def nest_get(dic, keys):
        result = dic
        for k in keys:
            result = result[k]
        return result
    

    Demo:

    >>> nest_get(my_dict, keys)
    >>> {'New Brunswick': ['Moncton', 'Saint John'], 'Saskatchewan': ['Saskatoon']}