Search code examples
pythonpython-3.xdictionarydefaultdict

Defaultdict that automatically creates a sub-dict for each key?


I have a dictionary where I am constantly doing stuff like this in my code:

special_dict = {}

# ...

if username not in special_dict:
    special_dict[username] = {}
    for subkey in ["Subkey1", "Subkey2", "Subkey3"]:
        special_dict[username][subkey] = []  # or {}, etc, depending on usecase

Basically I want a dictionary where for every username, the value is yet another dictionary of three specific subkeys, and then those values are lists or sets or what have you.

I'm familiar with defaultdict but I am not sure how to make the "value type" here something very specific. Normally I do defaultdict(list) if I want every value to be a list by default, but is there a way to make the default not a list but in itself a specific type of dictionary?

Ideally, in the end what I want to be able to do is special_dict[username][subkey].append(item) and not have to worry about whether or not the username exists first, because if it doesn't, it'll become a key and have the three subkeys formed.


Solution

  • You need a function that will create the structure you want, and pass this function as argument to defaultdict:

    from collections import defaultdict
    
    def name_subdict():
        return {'key1':[], 'key2':set(), 'key3':{}}
    
    mydict = defaultdict(name_subdict)
    
    mydict['John']['key1'].append(1)
    mydict['John']['key2'].add(2)
    mydict['Jane']['key3'][10] = 20
    
    print(mydict)
    # defaultdict(<function name_subdict at 0x7fcaf81193a0>, 
    # {'John': {'key1': [1], 'key2': {2}, 'key3': {}}, 
    #  'Jane': {'key1': [], 'key2': set(), 'key3': {10: 20}}})
    

    To answer your comment: yes, you can pass the type of data you want to be used for all subkeys, as in mydict = name_subdict(list). There is only one caveat: the argument to defaultdict must be a function (or any callable) that takes no argument.

    So, name_subdict(list) should return a function that will in turn create the structure.

    The code would then be:

    from collections import defaultdict
    
    def name_subdict(data_type):
        # data type must be a callable like list, set, dict...
        def subdict_creator():
            return {key:data_type() for key in ['key1', 'key2', 'key3']}
        return subdict_creator
        
    my_list_dict = defaultdict(name_subdict(list))
    my_set_dict = defaultdict(name_subdict(set))
    
    my_list_dict['John']['key1'].append(1)
    my_list_dict['John']['key2'].append(2)
    
    my_set_dict['Jane']['key3'].add(10)
    
    print(my_list_dict)
    # defaultdict(<function name_subdict.<locals>.subdict_creator at 0x7fcadbf27b80>, 
    # {'John': {'key1': [1], 'key2': [2], 'key3': []}})
    
    print(my_set_dict)
    # defaultdict(<function name_subdict.<locals>.subdict_creator at 0x7fcadbbf25e0>, 
    # {'Jane': {'key1': set(), 'key2': set(), 'key3': {10}}})