Toy recursive function that turns a list of keys and a value into a nested dict

list_nested_dict_toy.py

"""
This is a toy version of recursive function I wrote to update a value in an arbitrary
place in a nested dictionary.

I originally wrote something similar to this to handle writing to a TOML file
so that I could use a CLI tool such as
`my_tool set a.b.c.d hello` and it would write:
[a]
[a.b]
[a.b.c]
d = "hello"
"""

import copy


def get_value(d, keys):
    """ get_value takes a list of keys to find the value in a nested dictionary.
    it returns None if it reaches a dead end.
    example: 
        input: get_value({'a': {'b': {'c': 'hello'}}},[a,b,c])
        output: 'hello'
        
        input: get_value({'a': {'b': {'c': 'hello'}}},[a,b])
        output: {'c': 'hello'}
    """
    r = copy.deepcopy(d)  # keep it clean, kids.
    for k in keys:
        r = r.get(k)
        if not r:
            return None
    return r


def nested_value(table, keys, value, _depth=0):
    """ nested_value is a recursive function that takes a table (dict), list of keys,
    a value of any type, and internally uses a _depth counter.
    it safely traverses a dictionary in reverse order of the keys and either creates
    the data structures of updates the existing nested data structures if they arleady
    exist.

    example:
        input: nested_value({}, [a], 'foo')
        output: {'a': 'foo'}

        input: nested_value({}, [a,b,c,d], 'foo')
        output: {'a':{'b':{'c':{'d': 'foo'}}}}

        input: nested_value({'a':{'b':{'c':{'d': 'foo'}}}}, [a,b,c,e], 'bar')
        output: {'a':{'b':{'c':{'d': 'foo', 'e': 'bar'}}}}

    """
    assert len(keys) > 0
    
    k = keys.copy()  # we don't want to mutate, y'all.
    k.reverse()  # this is just a clever hack for depth=0
    if _depth == len(keys) - 1:
        current_val = get_value(table, keys)
        if isinstance(current_val, dict):
            raise TypeError(
                "Ovewriting a dict with another value is forbidden")
        table[keys[0]] = value
        return table

    parent_key = keys[:len(keys) - _depth - 1]
    parent_val = get_value(table, parent_key)

    if parent_val == None:
        parent_val = {}

    if not isinstance(parent_val, dict):
        raise TypeError(
            f"expecting: table, got: {type(parent_val)} for {parent_key}")

    parent_val[k[_depth]] = value
    return nested_value(table, keys, parent_val, _depth + 1)


test_table = [{
    'keys': ['a', 'b', 'c', 'd'],
    'value': 'hello',
    'table': {},
    'expected': {
        'a': {
            'b': {
                'c': {
                    'd': 'hello'
                }
            }
        }
    },
}, {
    'keys': ['a', 'b', 'c', 'e'],
    'value': 'world',
    'table': {
        'a': {
            'b': {
                'c': {
                    'd': 'hello'
                }
            }
        }
    },
    'expected': {
        'a': {
            'b': {
                'c': {
                    'd': 'hello',
                    'e': 'world',
                }
            }
        }
    }
}, {
    'keys': ['a'],
    'value': 'world',
    'table': {},
    'expected': {
        'a': "world"
    }
}]

# just running a few POC scenerios
for t in test_table:
    keys = t['keys']
    value = t['value']
    table = t['table']
    expected = t['expected']
    table = nested_value(table, keys, value)
    print(f'test:     {table}')
    print(f'expected: {expected}')
    assert table == expected