Improve `traverse_obj`

This commit is contained in:
pukkandan 2021-07-11 03:44:39 +05:30
parent 7dde84f3c9
commit 325ebc1703
No known key found for this signature in database
GPG Key ID: 0F00D95A001F4698
1 changed files with 34 additions and 22 deletions

View File

@ -6224,38 +6224,50 @@ def load_plugins(name, suffix, namespace):
return classes return classes
def traverse_obj(obj, keys, *, casesense=True, is_user_input=False, traverse_string=False): def traverse_obj(
obj, *key_list, default=None, expected_type=None,
casesense=True, is_user_input=False, traverse_string=False):
''' Traverse nested list/dict/tuple ''' Traverse nested list/dict/tuple
@param default Default value to return
@param expected_type Only accept final value of this type
@param casesense Whether to consider dictionary keys as case sensitive @param casesense Whether to consider dictionary keys as case sensitive
@param is_user_input Whether the keys are generated from user input. If True, @param is_user_input Whether the keys are generated from user input. If True,
strings are converted to int/slice if necessary strings are converted to int/slice if necessary
@param traverse_string Whether to traverse inside strings. If True, any @param traverse_string Whether to traverse inside strings. If True, any
non-compatible object will also be converted into a string non-compatible object will also be converted into a string
''' '''
keys = list(keys)[::-1]
while keys:
key = keys.pop()
if isinstance(obj, dict):
assert isinstance(key, compat_str)
if not casesense: if not casesense:
obj = {k.lower(): v for k, v in obj.items()} _lower = lambda k: k.lower() if isinstance(k, str) else k
key = key.lower() key_list = ((_lower(k) for k in keys) for keys in key_list)
obj = obj.get(key)
def _traverse_obj(obj, keys):
for key in list(keys):
if isinstance(obj, dict):
obj = (obj.get(key) if casesense or (key in obj)
else next((v for k, v in obj.items() if _lower(k) == key), None))
else: else:
if is_user_input: if is_user_input:
key = (int_or_none(key) if ':' not in key key = (int_or_none(key) if ':' not in key
else slice(*map(int_or_none, key.split(':')))) else slice(*map(int_or_none, key.split(':'))))
if key is None: if not isinstance(key, (int, slice)):
return None return None
if not isinstance(obj, (list, tuple)): if not isinstance(obj, (list, tuple)):
if traverse_string: if not traverse_string:
obj = compat_str(obj) return None
else: obj = str(obj)
try:
obj = obj[key]
except IndexError:
return None return None
assert isinstance(key, (int, slice))
obj = try_get(obj, lambda x: x[key])
return obj return obj
for keys in key_list:
val = _traverse_obj(obj, keys)
if val is not None:
if expected_type is None or isinstance(val, expected_type):
return val
return default
def traverse_dict(dictn, keys, casesense=True): def traverse_dict(dictn, keys, casesense=True):
''' For backward compatibility. Do not use ''' ''' For backward compatibility. Do not use '''