A convenient path-lookup script for python dictionaries

pathlike_dict_access.py

import sys, csv, itertools, json
#inspired by http://stackoverflow.com/questions/14692690/access-python-nested-dictionary-items-via-a-list-of-keys

# Testing dictionary
D = {
  "a": {
    "aa": {
      "aaa": 1,
      "aab": 2
    },
    "ab": {
      "aba": 3,
      "abb": 4
    }
  },
  "b": {
    "ba": 5,
    "bb": 6
  },
  "c": 7
}

# Some example paths
P1 = "a:aa:aaa" # 1
P2 = "a:aa:aab" # 2
P3 = "a:ab:aba" # 3
P4 = "a:ab" # {"aba": 3, "abb": 4}
P5 = "b:ba" # 5
P6 = "c" # 7
P7 = "a:ba:aba" # err

# returns the value at location specified by path
# and delimiter specified, defaults to :
def path_get(data_dict, path="", delim=":"):
  if path == "":
    return data_dict
  else:
    path_list = path.split(":")
    try:
      return reduce(lambda d, k: d[k], path_list, data_dict)
    except KeyError:
      print "[path_get] Innapropriate path: " + path

def path_set(data_dict, path, val="",  delim=":"):
  if delim not in path:
    data_dict[path] = val
    return path_get(data_dict, path, delim)
  else:
    path_list = path.rsplit(":",1)
    try:
      path_get(data_dict, path_list[0])[path_list[1]] = val
      return path_get(data_dict, path, delim)
    except KeyError:
      print "[path_set] Innapropriate destination path: " + path_list[0]

def main():
  assert (path_get(D) == D), "failed P0"
  assert (path_get(D,P1) == 1), "failed P1"
  assert (path_get(D,P2,":") == 2), "failed P2"
  assert (path_get(D,P3,":") == 3), "failed P3"
  assert (path_get(D,P4,":") == {'aba': 3, 'abb': 4}), "failed P4"
  assert (path_get(D,P5,":") == 5), "failed P5"
  assert (path_get(D,P6,":") == 7), "failed P6"
  # assert (path_get(D,P7,":") == 8), "failed P7"
  assert (path_set(D, "d", 10) == 10), "failed path set for 10"
  assert (path_set(D, "a:here", 11) == 11), "failed path set for 11"
  assert (path_set(D, "a:ab:new", 12) == 12), "failed path set for 12"
  # assert (path_set(D, "a:ba:new", 12) == 12), "failed path set for 12"

if __name__=="__main__":
  main()