From 325ebc1703beaa25074553242bb0a9b1399e699b Mon Sep 17 00:00:00 2001 From: pukkandan Date: Sun, 11 Jul 2021 03:44:39 +0530 Subject: [PATCH] Improve `traverse_obj` --- yt_dlp/utils.py | 56 ++++++++++++++++++++++++++++++------------------- 1 file changed, 34 insertions(+), 22 deletions(-) diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 888cfbb7e..8f9cb46f6 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -6224,37 +6224,49 @@ def load_plugins(name, suffix, namespace): 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 + @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 is_user_input Whether the keys are generated from user input. If True, strings are converted to int/slice if necessary @param traverse_string Whether to traverse inside strings. If True, any 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: - obj = {k.lower(): v for k, v in obj.items()} - key = key.lower() - obj = obj.get(key) - else: - if is_user_input: - key = (int_or_none(key) if ':' not in key - else slice(*map(int_or_none, key.split(':')))) - if key is None: + if not casesense: + _lower = lambda k: k.lower() if isinstance(k, str) else k + key_list = ((_lower(k) for k in keys) for keys in key_list) + + 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: + if is_user_input: + key = (int_or_none(key) if ':' not in key + else slice(*map(int_or_none, key.split(':')))) + if not isinstance(key, (int, slice)): return None - if not isinstance(obj, (list, tuple)): - if traverse_string: - obj = compat_str(obj) - else: + if not isinstance(obj, (list, tuple)): + if not traverse_string: + return None + obj = str(obj) + try: + obj = obj[key] + except IndexError: 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):