fix: improve deep_get, deep_set, deep_pop, and deep_iter functions for better error handling and performance

module/config/utils.py

@@ -2,6 +2,7 @@ import json
 import os
 import random
 import string
+from collections import deque
 from datetime import datetime, timedelta, timezone
 
 import yaml
@@ -190,57 +191,95 @@ def alas_instance():
 
 
 def deep_get(d, keys, default=None):
-    """
+    # 240 + 30 * depth (ns)
-    Get values in dictionary safely.
+    if type(keys) is str:
-    https://stackoverflow.com/questions/25833613/safe-method-to-get-value-of-nested-dictionary
-
-    Args:
-        d (dict):
-        keys (str, list): Such as `Scheduler.NextRun.value`
-        default: Default return if key not found.
-
-    Returns:
-
-    """
-    if isinstance(keys, str):
         keys = keys.split('.')
-    assert type(keys) is list
+
-    if d is None:
+    try:
-        return default
+        for k in keys:
-    if not keys:
+            d = d[k]
         return d
-    return deep_get(d.get(keys[0]), keys[1:], default)
+    # No such key
+    except KeyError:
+        return default
+    # Input `keys` is not iterable or input `d` is not dict
+    except TypeError:
+        return default
 
 
 def deep_set(d, keys, value):
     """
     Set value into dictionary safely, imitating deep_get().
     """
-    if isinstance(keys, str):
+    # 150 * depth (ns)
+    if type(keys) is str:
         keys = keys.split('.')
-    assert type(keys) is list
+
-    if not keys:
+    first = True
-        return value
+    exist = True
-    if not isinstance(d, dict):
+    prev_d = None
-        d = {}
+    prev_k = None
-    d[keys[0]] = deep_set(d.get(keys[0], {}), keys[1:], value)
+    prev_k2 = None
-    return d
+    try:
+        for k in keys:
+            if first:
+                prev_d = d
+                prev_k = k
+                first = False
+                continue
+            try:
+                # if key in dict: dict[key] > dict.get > dict.setdefault > try dict[key] except
+                if exist and prev_k in d:
+                    prev_d = d
+                    d = d[prev_k]
+                else:
+                    exist = False
+                    new = {}
+                    d[prev_k] = new
+                    d = new
+            except TypeError:
+                # `d` is not dict
+                exist = False
+                d = {}
+                prev_d[prev_k2] = {prev_k: d}
+
+            prev_k2 = prev_k
+            prev_k = k
+            # prev_k2, prev_k = prev_k, k
+    # Input `keys` is not iterable
+    except TypeError:
+        return
+
+    # Last key, set value
+    try:
+        d[prev_k] = value
+        return
+    # Last value `d` is not dict
+    except TypeError:
+        prev_d[prev_k2] = {prev_k: value}
+        return
 
 
 def deep_pop(d, keys, default=None):
-    """
+    if type(keys) is str:
-    Pop value from dictionary safely, imitating deep_get().
-    """
-    if isinstance(keys, str):
         keys = keys.split('.')
-    assert type(keys) is list
+
-    if not isinstance(d, dict):
+    try:
+        for k in keys[:-1]:
+            d = d[k]
+        return d.pop(keys[-1], default)
+    # No such key
+    except KeyError:
         return default
-    if not keys:
+    # Input `keys` is not iterable or input `d` is not dict
+    except TypeError:
+        return default
+    # Input `keys` out of index
+    except IndexError:
+        return default
+    # Last `d` is not dict
+    except AttributeError:
         return default
-    elif len(keys) == 1:
-        return d.pop(keys[0], default)
-    return deep_pop(d.get(keys[0]), keys[1:], default)
 
 
 def deep_default(d, keys, value):
@@ -262,26 +301,75 @@ def deep_default(d, keys, value):
     return d
 
 
-def deep_iter(data, depth=0, current_depth=1):
+def deep_iter(data, min_depth=None, depth=3):
     """
-    Iter a dictionary safely.
+    300us on alas.json depth=3 (530+ rows)
 
     Args:
-        data (dict):
+        data:
-        depth (int): Maximum depth to iter
+        min_depth:
-        current_depth (int):
+        depth:
 
     Returns:
-        list: Key path
+
-        Any:
     """
-    if isinstance(data, dict) \
+    if min_depth is None:
-            and (depth and current_depth <= depth):
+        min_depth = depth
-        for key, value in data.items():
+    assert 1 <= min_depth <= depth
-            for child_path, child_value in deep_iter(value, depth=depth, current_depth=current_depth + 1):
+
-                yield [key] + child_path, child_value
+    # Equivalent to dict.items()
-    else:
+    try:
-        yield [], data
+        if depth == 1:
+            for k, v in data.items():
+                yield [k], v
+            return
+        # Iter first depth
+        elif min_depth == 1:
+            q = deque()
+            for k, v in data.items():
+                key = [k]
+                if type(v) is dict:
+                    q.append((key, v))
+                else:
+                    yield key, v
+        # Iter target depth only
+        else:
+            q = deque()
+            for k, v in data.items():
+                key = [k]
+                if type(v) is dict:
+                    q.append((key, v))
+    except AttributeError:
+        # `data` is not dict
+        return
+
+    # Iter depths
+    current = 2
+    while current <= depth:
+        new_q = deque()
+        # max depth
+        if current == depth:
+            for key, data in q:
+                for k, v in data.items():
+                    yield key + [k], v
+        # in target depth
+        elif min_depth <= current < depth:
+            for key, data in q:
+                for k, v in data.items():
+                    subkey = key + [k]
+                    if type(v) is dict:
+                        new_q.append((subkey, v))
+                    else:
+                        yield subkey, v
+        # Haven't reached min depth
+        else:
+            for key, data in q:
+                for k, v in data.items():
+                    subkey = key + [k]
+                    if type(v) is dict:
+                        new_q.append((subkey, v))
+        q = new_q
+        current += 1
 
 
 def parse_value(value, data):