mirror of
https://github.com/yt-dlp/yt-dlp.git
synced 2025-01-19 09:37:34 +01:00
[utils] Improve traverse_obj
* Allow skipping a level: `traverse_obj([{k:v1}, {k:v2}], (None, k))` => `[v1, v2]` * Make keys variadic: `traverse_obj(obj, k1: str, k2: str)` => `traverse_obj(obj, (k1,), (k2,))` * Fetch from multiple keys: `traverse_obj([{k1:[1], k2:[2], k3:[3]}], (0, (k1, k2), 0))` => `[1, 2]` TODO: Add tests
This commit is contained in:
parent
0ba692acc8
commit
8f3343809e
@ -6225,9 +6225,14 @@ def load_plugins(name, suffix, namespace):
|
|||||||
|
|
||||||
|
|
||||||
def traverse_obj(
|
def traverse_obj(
|
||||||
obj, *key_list, default=None, expected_type=None,
|
obj, *path_list, default=None, expected_type=None,
|
||||||
casesense=True, is_user_input=False, traverse_string=False):
|
casesense=True, is_user_input=False, traverse_string=False):
|
||||||
''' Traverse nested list/dict/tuple
|
''' 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,
|
||||||
|
all the keys given in the tuple are traversed, and
|
||||||
|
"..." traverses all the keys in the object
|
||||||
@param default Default value to return
|
@param default Default value to return
|
||||||
@param expected_type Only accept final value of this type
|
@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
|
||||||
@ -6235,23 +6240,38 @@ def traverse_obj(
|
|||||||
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
|
||||||
|
# TODO: Write tests
|
||||||
'''
|
'''
|
||||||
if not casesense:
|
if not casesense:
|
||||||
_lower = lambda k: k.lower() if isinstance(k, str) else k
|
_lower = lambda k: k.lower() if isinstance(k, str) else k
|
||||||
key_list = ((_lower(k) for k in keys) for keys in key_list)
|
path_list = (map(_lower, variadic(path)) for path in path_list)
|
||||||
|
|
||||||
def _traverse_obj(obj, keys):
|
def _traverse_obj(obj, path, _current_depth=0):
|
||||||
for key in list(keys):
|
nonlocal depth
|
||||||
if isinstance(obj, dict):
|
path = tuple(variadic(path))
|
||||||
|
for i, key in enumerate(path):
|
||||||
|
if isinstance(key, (list, tuple)):
|
||||||
|
obj = [_traverse_obj(obj, sub_key, _current_depth) for sub_key in key]
|
||||||
|
key = ...
|
||||||
|
if key is ...:
|
||||||
|
obj = (obj.values() if isinstance(obj, dict)
|
||||||
|
else obj if isinstance(obj, (list, tuple, LazyList))
|
||||||
|
else str(obj) if traverse_string else [])
|
||||||
|
_current_depth += 1
|
||||||
|
depth = max(depth, _current_depth)
|
||||||
|
return [_traverse_obj(inner_obj, path[i + 1:], _current_depth) for inner_obj in obj]
|
||||||
|
elif isinstance(obj, dict):
|
||||||
obj = (obj.get(key) if casesense or (key in obj)
|
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 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 == slice(None):
|
||||||
|
return _traverse_obj(obj, (..., *path[i + 1:]))
|
||||||
if not isinstance(key, (int, slice)):
|
if not isinstance(key, (int, slice)):
|
||||||
return None
|
return None
|
||||||
if not isinstance(obj, (list, tuple)):
|
if not isinstance(obj, (list, tuple, LazyList)):
|
||||||
if not traverse_string:
|
if not traverse_string:
|
||||||
return None
|
return None
|
||||||
obj = str(obj)
|
obj = str(obj)
|
||||||
@ -6261,10 +6281,18 @@ def traverse_obj(
|
|||||||
return None
|
return None
|
||||||
return obj
|
return obj
|
||||||
|
|
||||||
for keys in key_list:
|
for path in path_list:
|
||||||
val = _traverse_obj(obj, keys)
|
depth = 0
|
||||||
|
val = _traverse_obj(obj, path)
|
||||||
if val is not None:
|
if val is not None:
|
||||||
if expected_type is None or isinstance(val, expected_type):
|
if depth:
|
||||||
|
for _ in range(depth - 1):
|
||||||
|
val = itertools.chain.from_iterable(filter(None, val))
|
||||||
|
val = (list(filter(None, val)) if expected_type is None
|
||||||
|
else [v for v in val if isinstance(v, expected_type)])
|
||||||
|
if val:
|
||||||
|
return val
|
||||||
|
elif expected_type is None or isinstance(val, expected_type):
|
||||||
return val
|
return val
|
||||||
return default
|
return default
|
||||||
|
|
||||||
|
Loading…
x
Reference in New Issue
Block a user