norbertpy norbertpy - 1 month ago 23
Python Question

Python nested defaultdict with mix data types

So, how can I create a defaultdict for this:

{
'branch': {
'count': 23,
'leaf': {
'tag1': 30,
'tag2': 10
}
},
}


so that, I'll get zeros for
count
,
tag1
and
tag2
as default? I wanna populate the dict dynamically while I'm reading the inputs. When I see a new
branch
I want to create a dict with
count
as zero and an empty dict as leaf. When I get a
leaf
, I want to create a key with it's name and set the value to zero.

Update:
Accepted Martijn's answer as it has more upvotes but other answers are equally good.

Answer

You can't do this with defaultdict, because the factory doesn't have access to the key.

However, you can just subclass dict to create your own 'smart' defaultdict-like class. Provide your own __missing__ method that adds values based on the key:

class KeyBasedDefaultDict(dict):
    def __init__(self, default_factories, *args, **kw):
        self._default_factories = default_factories
        super(KeyBasedDefaultDict, self).__init__(*args, **kw)

    def __missing__(self, key):
        factory = self._default_factories.get(key)
        if factory is None:
            raise KeyError(key)
        new_value = factory()
        self[key] = new_value
        return new_value

Now you can supply your own mapping:

mapping = {'count': int, 'leaf': dict}
mapping['branch'] = lambda: KeyBasedDefaultDict(mapping)

tree = KeyBasedDefaultDict(mapping)

Demo:

>>> mapping = {'count': int, 'leaf': dict}
>>> mapping['branch'] = lambda: KeyBasedDefaultDict(mapping)
>>> tree = KeyBasedDefaultDict(mapping)
>>> tree['branch']['count'] += 23
>>> tree['branch']['leaf']['tag1'] = 30
>>> tree['branch']['leaf']['tag2'] = 10
>>> tree
{'branch': {'count': 23, 'leaf': {'tag1': 30, 'tag2': 10}}}