From 2614f64600f9249682897786f5345a61d98dafeb Mon Sep 17 00:00:00 2001 From: pukkandan Date: Sat, 9 Oct 2021 08:14:41 +0530 Subject: [PATCH] [utils] Let traverse_obj accept functions as keys --- yt_dlp/utils.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 027387897..db9b9de94 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -6335,7 +6335,9 @@ def traverse_obj( ''' Traverse nested list/dict/tuple @param path_list A list of paths which are checked one by one. Each path is a list of keys where each key is a string, - a tuple of strings or "...". When a tuple is given, + a function, a tuple of strings or "...". + When a fuction is given, it takes the key as argument and + returns whether the key matches or not. When a tuple is given, all the keys given in the tuple are traversed, and "..." traverses all the keys in the object @param default Default value to return @@ -6368,6 +6370,18 @@ def traverse_obj( _current_depth += 1 depth = max(depth, _current_depth) return [_traverse_obj(inner_obj, path[i + 1:], _current_depth) for inner_obj in obj] + elif callable(key): + if isinstance(obj, (list, tuple, LazyList)): + obj = enumerate(obj) + elif isinstance(obj, dict): + obj = obj.items() + else: + if not traverse_string: + return None + obj = str(obj) + _current_depth += 1 + depth = max(depth, _current_depth) + return [_traverse_obj(v, path[i + 1:], _current_depth) for k, v in obj if key(k)] elif isinstance(obj, dict) and not (is_user_input and key == ':'): obj = (obj.get(key) if casesense or (key in obj) else next((v for k, v in obj.items() if _lower(k) == key), None))