dot_paths.py

import re


class _ValueNotSet(object):
    pass


ValueNotSet = _ValueNotSet()


def get_by_dot_path(container, path, default_value=ValueNotSet):
    if not isinstance(path, basestring):
        raise ValueError('Path must be a string')

    if container is not None and not hasattr(container, '__getitem__'):
        raise ValueError('Container must support __getitem__')

    try:
        if not container:
            raise KeyError(path)

        if '.' not in path and '[' not in path and not re.match('^[0-9]+$', path):
            return container[path]

        path_parts = get_dot_path_parts(path)

        current = container
        while path_parts:
            if not current:
                raise KeyError(path)

            if not isinstance(current, (dict, list, tuple)):
                raise KeyError(path)

            part = path_parts.pop(0)
            if re.match('^[0-9]+$', part) and str(int(part)) == part:
                int_part = int(part)
            else:
                int_part = None

            if isinstance(current, (list, tuple)):
                if int_part is None:
                    raise KeyError(path)
                if len(current) <= int_part:
                    raise KeyError(path)
                current = current[int_part]
            else:
                if part in current:
                    current = current[part]
                elif int_part in current:
                    current = current[int_part]
                else:
                    raise KeyError(path)

            if not path_parts:
                return current

    except KeyError:
        if default_value is ValueNotSet:
            raise
        else:
            return default_value


def get_dot_path_parts(path):
    path_parts = [p for p in re.split('[\.\[\]]+', path) if p]
    return path_parts


def exists_dot_path(container, path):
    try:
        get_by_dot_path(container, path)
        return True
    except KeyError:
        return False


def set_by_dot_path(container, path, value, full_path=None):
    """
    This requires the container path to exist because we don't want to be creating the wrong kind of container
    (list or dict).
    """
    assert isinstance(path, basestring)
    assert isinstance(container, (list, dict))

    if full_path is None:
        full_path = path

    path_parts = get_dot_path_parts(path)
    if len(path_parts) == 1:
        part = path_parts[0]
        int_part = None
        if re.match('^[0-9]+$', path) and str(int(part)) == str(part):
            int_part = int(path)

        if isinstance(container, list):
            if int_part is not None:
                container[int_part] = value
            else:
                raise ValueError(path)
        else:
            if part in container:
                container[part] = value
            elif int_part in container:
                container[int_part] = value
            else:
                container[part] = value
    else:
        rest_of_path = '.'.join(path_parts[1:])

        part = path_parts[0]
        int_part = None
        if re.match('^[0-9]+$', path) and str(int(part)) == str(part):
            int_part = int(path)

        if isinstance(container, list):
            if int_part is not None:
                assert exists_dot_path(container[int_part], rest_of_path)
                set_by_dot_path(container[int_part], rest_of_path, value, full_path=full_path)
            else:
                raise ValueError(path)
        else:
            if part in container:
                assert exists_dot_path(container[part], rest_of_path)
                set_by_dot_path(container[part], rest_of_path, value, full_path=full_path)
            elif int_part in container:
                assert exists_dot_path(container[int_part], rest_of_path)
                set_by_dot_path(container[int_part], rest_of_path, value, full_path=full_path)
            else:
                raise ValueError(path)


if __name__ == '__main__':
    list_obj = ['a', 'b', 'c', {'d': {'one': 1, 'two': 2, 'three': 3}}]
    assert get_by_dot_path(list_obj, '0') == 'a'
    assert get_by_dot_path(list_obj, '[0]') == 'a'
    assert get_by_dot_path(list_obj, '[3].d')
    assert get_by_dot_path(list_obj, '[3].d.two') == 2
    assert not exists_dot_path(list_obj, '[3].d.four')

    dict_obj = {'a': {'one': 1, 'two': 2, 'three': {'x': 1, 'y': 2}}, 'b': ['x', 'y', 'z'], 'c': 'ccc'}
    assert get_by_dot_path(dict_obj, 'a')
    assert not exists_dot_path(dict_obj, 'a[0]')
    assert get_by_dot_path(dict_obj, 'a.one')
    assert get_by_dot_path(dict_obj, 'a.three')
    assert get_by_dot_path(dict_obj, 'a.three.x')
    assert not exists_dot_path(dict_obj, 'a.three.x.xxx')
    assert not exists_dot_path(dict_obj, 'a.three.x[1]')
    assert get_by_dot_path(dict_obj, 'b[0]') == 'x'
    assert get_by_dot_path(dict_obj, 'b[2]') == 'z'
    assert not exists_dot_path(dict_obj, 'b[3]')
    assert get_by_dot_path(dict_obj, 'c') == 'ccc'
    assert not exists_dot_path(dict_obj, 'c[0]')

    set_by_dot_path(dict_obj, 'c', 'c4')
    assert get_by_dot_path(dict_obj, 'c') == 'c4'

    set_by_dot_path(dict_obj, 'a.three.x', {'xxx': 'yyy'})
    assert get_by_dot_path(dict_obj, 'a.three.x.xxx') == 'yyy'

    set_by_dot_path(dict_obj, '_id', 23)
    assert get_by_dot_path(dict_obj, '_id') == 23

    set_by_dot_path(list_obj, '1', 'hahaha')
    assert get_by_dot_path(list_obj, '1') == 'hahaha'