config_dict: T2668: move keyword arg get_first_key into get_sub_dict

2 files changed, 12 insertions, 8 deletions

diff --git a/python/vyos/config.py b/python/vyos/config.py
index 4ebb68a09..6ad45819a 100644
--- a/python/vyos/config.py
+++ b/python/vyos/config.py
@@ -299,12 +299,7 @@ class Config(object):
         else:
             config_tree = vyos.configtree.ConfigTree(res)
             config_dict = json.loads(config_tree.to_json())
-            config_dict = vyos.util.get_sub_dict(config_dict, self._make_path(path))
-            if get_first_key and path and config_dict:
-                tmp = next(iter(config_dict.values()))
-                if not isinstance(tmp, dict):
-                    raise TypeError("Data under node is not of type dict")
-                config_dict = tmp
+            config_dict = vyos.util.get_sub_dict(config_dict, self._make_path(path), get_first_key)
 
         if key_mangling:
             if not (isinstance(key_mangling, tuple) and \
diff --git a/python/vyos/util.py b/python/vyos/util.py
index 9b6b26a3b..924df6b3a 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -376,7 +376,7 @@ def _get_sub_dict(d, lpath):
         return {}
     return _get_sub_dict(c[k], lpath)
 
-def get_sub_dict(source, lpath):
+def get_sub_dict(source, lpath, get_first_key=False):
     """ Returns the sub-dict of a nested dict, defined by path of keys.
 
     Args:
@@ -393,7 +393,16 @@ def get_sub_dict(source, lpath):
         raise TypeError("path must be of type list")
     if not lpath:
         return source
-    return _get_sub_dict(source, lpath)
+
+    ret =  _get_sub_dict(source, lpath)
+
+    if get_first_key and lpath and ret:
+        tmp = next(iter(ret.values()))
+        if not isinstance(tmp, dict):
+            raise TypeError("Data under node is not of type dict")
+        ret = tmp
+
+    return ret
 
 def process_running(pid_file):
     """ Checks if a process with PID in pid_file is running """