Python 2 - Searching for all the values of a key in a dict with nested dicts and lists

Question

I've got a dictionary with nested dicts and list inside it. This structure allows a key to be repeated inside the dict, being able to have different values (either from a list of dicts, or from different levels of the dict). My solution use recursion in order to search on the keys of the dict, but also on the dicts or lists that may be nested inside it. The thing is that, if there is a multiple occurence of a key at different level, the output of the method will only return the higher level appearance.

I've tried using yield, but I've struggled to do it and got mental-boomed.

def search_dict_list(collection, desired_key):
    """Returns the value of every element with a desired_key from a response 
    dictionary which may contain nested lists and dictionaries.
    If it doesn't exist, it returns None

    :param collection: response body dictionary
    :type collection: dict
    :param desired_key: key of the elements wanted to find
    :type desired_key: str
    """
    if isinstance(collection, dict):
        if desired_key in list(collection.keys()):
            item = search_dict_list(collection[desired_key], desired_key)
            return [collection[desired_key], item] if item else collection[desired_key]
        else:
            for element in list(collection.values()):
                item = search_dict_list(element, desired_key)
                if item is not None:
                    return item
    if isinstance(collection, list):
        for element in collection:
            item = search_dict_list(element, desired_key)
            if item is not None:
                return item

test_dict = {
    "mock_key_1": [
        {
            "mock_key_2": [
                {
                    "mock_key_3": {
                        "mock_key_4": {
                            "mock_key_5": [
                                {"target_key": "target_value_1"},
                                {"target_key": "target_value_2"},
                            ],
                        },
                        "mock_key_6": {
                            "mock_key_7": "mock_value_7",
                            "mock_key_8": "mock_value_8",
                            "mock_key_9": "mock_value_9",
                        },
                        "target_key": ["target_value_3"],
                    }
                }
            ],
            "target_key": {"mock_key_10" : "target_value_4"},
        }
    ]
}

The expected output here would include

[{"mock_key_10" : "target_value_4"}, ["target_value_3"], "target_value_1", "target_value_2"]

Answer 1

You can recursive function like this

def find_key_values(desired_key, collection):
    result = []

    if isinstance(collection, dict):
        # If the desired key is present in the dictionary,
        # add the corresponding value to the result list
        if desired_key in collection:
            result.append(collection[desired_key])

        for value in collection.values():
            result.extend(find_key_values(desired_key, value))

    elif isinstance(collection, list):
        for item in collection:
            result.extend(find_key_values(desired_key, item))

    return result


results = find_key_values('target_key', test_dict)

The results.extend() method is used to append the values returned by the recursive call to the find_key_values function to the results list.

This ensures that all the values associated with the key in the nested structure are included in the final results list.

Answer 2

You can use the following:

def search_dict_list(collection, desired_key):
    if isinstance(collection, dict):
        for k, v in collection.items():
            if k == desired_key:
                yield v
            yield from search_dict_list(v, desired_key)
    elif isinstance(collection, list):
        for item in collection:
            yield from search_dict_list(item, desired_key)


test_dict = {
    "mock_key_1": [
        {
            "mock_key_2": [
                {
                    "mock_key_3": {
                        "mock_key_4": {
                            "mock_key_5": [
                                {"target_key": "target_value_1"},
                                {"target_key": "target_value_2"},
                            ],
                        },
                        "mock_key_6": {
                            "mock_key_7": "mock_value_7",
                            "mock_key_8": "mock_value_8",
                            "mock_key_9": "mock_value_9",
                        },
                        "target_key": ["target_value_3"],
                    }
                }
            ],
            "target_key": {"mock_key_10": "target_value_4"},
        }
    ]
}

print(list(search_dict_list(test_dict, "target_key")))

output:

['target_value_1', 'target_value_2', ['target_value_3'], {'mock_key_10': 'target_value_4'}]

Basically it checks the type of the collection, if it's a list it delegates the checking to the individual items. If it's a dictionary, it checks the keys, if the desired key is found it returns its value.