Python nested defaultdict with mix data types

Question

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 1

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

Answer 2

An object has a __dict__ that stores data, and allows you to programmatically set defaults. There's also an object called a Counter which I think you should use to delegate the counting of your leafs.

Thus, I recommend that you use an object with a collections.Counter:

import collections

class Branch(object):
    def __init__(self, leafs=(), count=0):
        self.leafs = collections.Counter(leafs)
        self.count = count
    def __repr__(self):
        return 'Branch(leafs={0}, count={1})'.format(self.leafs, self.count)

BRANCHES = [Branch(['leaf1', 'leaf2']),
            Branch(['leaf3', 'leaf4', 'leaf3']),
            Branch(['leaf6', 'leaf7']),
           ]

And usage:

>>> import pprint
>>> pprint.pprint(BRANCHES)
[Branch(leafs=Counter({'leaf1': 1, 'leaf2': 1}), count=0),
 Branch(leafs=Counter({'leaf3': 2, 'leaf4': 1}), count=0),
 Branch(leafs=Counter({'leaf7': 1, 'leaf6': 1}), count=0)]
>>> first_branch = BRANCHES[0]
>>> first_branch.count += 23
>>> first_branch
Branch(leafs=Counter({'leaf1': 1, 'leaf2': 1}), count=23)
>>> first_branch.leafs['leaf that does not exist']
0
>>> first_branch.leafs.update(['new leaf'])
>>> first_branch
Branch(leafs=Counter({'new leaf': 1, 'leaf1': 1, 'leaf2': 1}), count=23)