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

Question

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.

Answer 1

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'}

Answer 2

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)