Merge pull request #486 from jestabro/sub_dict

2 files changed, 42 insertions, 6 deletions

diff --git a/python/vyos/config.py b/python/vyos/config.py
index 56353c322..0bbfece6e 100644
--- a/python/vyos/config.py
+++ b/python/vyos/config.py
@@ -71,7 +71,6 @@ import subprocess
 import vyos.util
 import vyos.configtree
 
-
 class VyOSError(Exception):
     """
     Raised on config access errors, most commonly if the type of a config tree node
@@ -288,17 +287,24 @@ class Config(object):
             self.__session_env = save_env
             return(default)
 
-    def get_config_dict(self, path=[], effective=False, key_mangling=None):
+    def get_config_dict(self, path=[], effective=False, key_mangling=None, get_first_key=False):
         """
         Args: path (str list): Configuration tree path, can be empty
         Returns: a dict representation of the config
         """
-        res = self.show_config(self._make_path(path), effective=effective)
-        if res:
+        res = self.show_config(effective=effective)
+        config_dict = {}
+        if not res:
+            return config_dict
+        else:
             config_tree = vyos.configtree.ConfigTree(res)
             config_dict = json.loads(config_tree.to_json())
-        else:
-            config_dict = {}
+            config_dict = vyos.util.get_sub_dict(config_dict, self._make_path(path))
+            if get_first_key:
+                tmp = next(iter(config_dict.values()))
+                if not isinstance(tmp, dict):
+                    raise TypeError("Data under node is not of type dict")
+                config_dict = tmp
 
         if key_mangling:
             if not (isinstance(key_mangling, tuple) and \
diff --git a/python/vyos/util.py b/python/vyos/util.py
index 0ddc14963..268eaa69a 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -364,6 +364,36 @@ def mangle_dict_keys(data, regex, replacement):
 
     return new_dict
 
+def _get_sub_dict(d, lpath):
+    k = lpath[0]
+    if k not in d.keys():
+        return {}
+    c = {k: d[k]}
+    lpath = lpath[1:]
+    if not lpath:
+        return c
+    elif not isinstance(c[k], dict):
+        return {}
+    return _get_sub_dict(c[k], lpath)
+
+def get_sub_dict(source, lpath):
+    """ Returns the sub-dict of a nested dict, defined by path of keys.
+
+    Args:
+        source (dict): Source dict to extract from
+        lpath (list[str]): sequence of keys
+
+    Returns: {key : source[..]..[key]} for key the last element of lpath, if exists
+             {} otherwise
+    """
+    if not isinstance(source, dict):
+        raise TypeError("source must be of type dict")
+    if not isinstance(lpath, list):
+        raise TypeError("path must be of type list")
+    if not lpath:
+        return source
+    return _get_sub_dict(source, lpath)
+
 def process_running(pid_file):
     """ Checks if a process with PID in pid_file is running """
     from psutil import pid_exists