writing and accessing nested dictionary values - python 2 and 3 compatible

dicts.py

import collections
import itertools
import operator

import six


def lenient_getitem(obj, key):
    try:
        return operator.getitem(obj, key)
    except KeyError:
        obj[key] = {}
        return obj[key]


def get_recursive_defaultdict():
    """Return a new defaultdict of defaultdicts."""
    def _ddf():
        return collections.defaultdict(_ddf)
    return collections.defaultdict(_ddf)


def write_nested_value(keys, value, initial=None, sep=None, strict=True):
    """Write (nested) values into a dictionary."""
    keys = keys if not sep else keys.split(sep)
    dicts = (dict, collections.Mapping)
    if not initial and not isinstance(initial, dicts):
        initial = get_recursive_defaultdict()
    getitem = operator.getitem if strict else lenient_getitem
    # play nicely with all kinds of sequences and iterables
    sequence = itertools.islice(keys, None, len(keys) - 1)
    fin = six.moves.reduce(getitem, sequence, initial)
    fin.update({keys[len(keys) - 1]: value})
    return initial


def read_nested_value(dictionary, keys, sep=None):
    """Return the value pointed to by sequence 'keys'."""
    keys = keys if not sep else keys.split(sep)
    return six.moves.reduce(operator.getitem, keys, dictionary)

ex.py

import dicts


# 1, test creates a dict on the fly
d = dicts.write_nested_value([1, 2, 3], 'THREE')
assert d[1][2][3] == 'THREE'
assert dicts.read_nested_value(d, [1, 2, 3]) == 'THREE'
del d

# 2, (python 2 or 3), list OR immutable Sequence, respectively
d = {1: {2: {3: {}}}}
r = dicts.write_nested_value(range(1, 5), 'FOUR', d)
assert r is d
assert d[1][2][3][4] == 'FOUR'
assert dicts.read_nested_value(d, [1, 2, 3]) == {4: 'FOUR'}
del d

# 3, test overwrites value
d = {1: {2: {'foo': 'bar'}}}
r = dicts.write_nested_value([1, 2, 'foo'], 'baz', d)
assert r is d
assert d[1][2]['foo'] == 'baz'
del d

# 4, test dont clobber, update the nested dict
d = {1: {2: {'foo': 'bar'}}}
r = dicts.write_nested_value([1, 2, 'zip'], 'bing', d)
assert r is d
assert d[1][2] == {'foo': 'bar', 'zip': 'bing'}
del d

# 5, test string as an iterable
d = dicts.write_nested_value('yes', 'ok')
assert d['y']['e']['s'] == 'ok'
assert dicts.read_nested_value(d, 'yes') == 'ok'
del d

# 6, test custom separator
d = dicts.write_nested_value('y/e/s', 'ok', sep='/')
assert d['y']['e']['s'] == 'ok'
assert dicts.read_nested_value(d, 'y/e/s', sep='/') == 'ok'
del d

# 7, test strict=False
d = {}
dicts.write_nested_value('abc', '+1', d, strict=False)
assert d == {'a': {'b': {'c': '+1'}}}
del d

# 8, test assert raises
d = {}
err = None
try:
    # strict=True
    dicts.write_nested_value('abc', '+1', d)
except KeyError as e:
    err = e
    pass
assert isinstance(err, KeyError)
assert err.args == ('a',)
del d