pick data from complex json by key path string

pick_from_json.py

# coding: utf-8
__author__ = 'Jux.Liu'


def pick_from_json(json_obj, key_path, key_sep='.', list_index_start='$_', list_index_end='_$'):
    if key_sep in list_index_start or key_path in list_index_end:
        raise Exception('pick another key sep')

    key_list = key_path.split(key_sep)
    sub_obj = json_obj
    itered = []
    for key in key_list:
        itered.append(key)
        if isinstance(sub_obj, dict):
            if key in sub_obj:
                sub_obj = sub_obj.get(key)
            else:
                raise KeyError(
                    'key: <{}> not exists, please check your key path. failed at: <{}>'.format(key, '.'.join(itered)))
        elif isinstance(sub_obj, list):
            index = int(key.lstrip(list_index_start).rstrip(list_index_end))
            if len(sub_obj) > index:
                sub_obj = sub_obj[index]
            else:
                raise IndexError('list index out of range. failed at: <{}>'.format('.'.join(itered)))
    return sub_obj


json_obj = {'a': ['foo1', 'foo2', 'foo3'], 'b': 123, 'c': {'qq': '123', 'pp': [{'test': 'bar'}, {'prod': 'par'}]}, 'd': None}
key_path_success = 'c.pp.1.prod'
key_path_fail = 'c.pp.3.test'

print(pick_from_json(json_obj, key_path_success))
print(pick_from_json(json_obj, key_path_fail))