How to mask a Python 3 nested dictionary to return a new dictionary with only certain items?

user_3268298 0
Viewed 4055

I am consuming several endpoints of an API that is very verbose in the data it returns. I would like to provide a subset of this data to another piece of code elsewhere.

Suppose I am given several dictionaries like this (which I plan to loop through and filter):

asset = {
    'id': 1,
    'name': 'MY-PC',
    'owner': 'me',
    'location': 'New York City',
    'model': {
        'id': 1,
        'name': 'Surface',
        'manufacturer': {
            'id': 1,
            'name': 'Microsoft'
        }
    }
}

I want to create a function that will take that dictionary in, along with a "mask" which will be used to create a new dictionary of only the allowed items. This might be an example mask (though, I can work with whatever format makes the resulting code the most concise):

mask = {
    'id': True,
    'name': True,
    'model': {
        'id': True,
        'name': True,
        'manufacturer': {
            'name': True
        }
    }
}

The function should then return this:

mask = {
    'id': 1,
    'name': 'MY-PC',
    'model': {
        'id': 1,
        'name': 'Surface',
        'manufacturer': {
            'name': 'Microsoft'
        }
    }
}

Is there something already built into Python 3 that would help aid in this? It looks like if I have to do this manually, it's going to get quite ugly quickly. I found itertools.compress, but that seems like it's for lists and won't handle the complexity of dictionaries.

iterator dictionary python
3 Answers

You can recursively build a new dict from the mask by selecting only values corresponding in the main dict:

def prune_dict(dct, mask):
    result = {}
    for k, v in mask.items():
        if isinstance(v, dict):
            value = prune_dict(dct[k], v)
            if value: # check that dict is non-empty
                result[k] = value
        elif v:
            result[k] = dct[k]
    return result

print(prune_dict(asset, mask))

{'id': 1,
'model': {'id': 1, 'manufacturer': {'name': 'Microsoft'}, 'name': 'Surface'},
'name': 'MY-PC'}
user_31255660

This would be a good chance to use recursion, here is some sample code I haven't tested:

def copy(asset, result, mask):
    for key_name, value in mask.items():
        if value == True:
            result[key_name] = asset[key_name]
        else:
            result[key_name] = x = {}
            copy(asset[key_name], x, value)

y = {}
copy(asset, y, mask)
user_6320880

This would probably be a recursive function. Also, for the mask, I recommend this format: mask = ["id", "name", "model.id", "model.name", "model.manufacturer.name"]

Then, you'd first keep only entries that are named in the mask:

def filterstage1(dictionary, mask):
    result = {}
    for key in dictionary:
        if isinstance(dictionary[key], dict):
            newmask = [maskname[mask.find(".") + 1:] for maskname in mask if maskname.startswith(key + ".")]
            result[k] = filterstage1(dictionary[key], newmask)
        elif key in mask:
            result[key] = dictionary[key]
    return result

Then, depending on whether or not you want to remove sub-dictionaries that were not in the mask and had no subelements, you can include the second stage:

def filterstage2(dictionary, mask):
    result = {}
    for key in dictionary:
        if not (isinstance(dictionary[key], dict) and dictionary[key] == {} and key not in mask):
            result[key] = dictionary[key]

Final code: filterstage2(filterstage1(dictionary, mask), mask). You can combine the two stages together if you wish.

user_82004850
Related