tree_defaultdict.py

from typing import Dict, Any, DefaultDict, Optional, TYPE_CHECKING
from collections import defaultdict

if TYPE_CHECKING:
    TreeDict = Dict[Any, Union[Any, "TreeDict"]]
    TreeDefaultDict = DefaultDict[Any, Union[Any, "TreeDefaultDict"]]


def tree_dict_factory(nested_dict: Optional["TreeDict"] = None) -> "TreeDefaultDict":
    """
    recursive tree default dict.
    allows for:
    tree[1][2][3][4][5] = 0
    even if none of those dict levels existed before.
    Can take a existing nested dict as a parameter
    """
    if nested_dict:
        d = defaultdict(tree_dict_factory)
        for k, v in nested_dict.items():
            if isinstance(v, dict):
                d[k] = tree_dict_factory(v)

        return defaultdict(tree_dict_factory, d)
    else:
        return defaultdict(tree_dict_factory)