]>
jfr.im git - yt-dlp.git/blob - yt_dlp/utils/traversal.py
19 obj
, *paths
, default
=NO_DEFAULT
, expected_type
=None, get_all
=True,
20 casesense
=True, is_user_input
=False, traverse_string
=False):
22 Safely traverse nested `dict`s and `Iterable`s
24 >>> obj = [{}, {"key": "value"}]
25 >>> traverse_obj(obj, (1, "key"))
28 Each of the provided `paths` is tested and the first producing a valid result will be returned.
29 The next path will also be tested if the path branched but no results could be found.
30 Supported values for traversal are `Mapping`, `Iterable` and `re.Match`.
31 Unhelpful values (`{}`, `None`) are treated as the absence of a value and discarded.
33 The paths will be wrapped in `variadic`, so that `'key'` is conveniently the same as `('key', )`.
35 The keys in the path can be one of:
36 - `None`: Return the current object.
37 - `set`: Requires the only item in the set to be a type or function,
38 like `{type}`/`{func}`. If a `type`, returns only values
39 of this type. If a function, returns `func(obj)`.
40 - `str`/`int`: Return `obj[key]`. For `re.Match`, return `obj.group(key)`.
41 - `slice`: Branch out and return all values in `obj[key]`.
42 - `Ellipsis`: Branch out and return a list of all values.
43 - `tuple`/`list`: Branch out and return a list of all matching values.
44 Read as: `[traverse_obj(obj, branch) for branch in branches]`.
45 - `function`: Branch out and return values filtered by the function.
46 Read as: `[value for key, value in obj if function(key, value)]`.
47 For `Iterable`s, `key` is the index of the value.
48 For `re.Match`es, `key` is the group number (0 = full match)
49 as well as additionally any group names, if given.
50 - `dict` Transform the current object and return a matching dict.
51 Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
53 `tuple`, `list`, and `dict` all support nested paths and branches.
55 @params paths Paths which to traverse by.
56 @param default Value to return if the paths do not match.
57 If the last key in the path is a `dict`, it will apply to each value inside
58 the dict instead, depth first. Try to avoid if using nested `dict` keys.
59 @param expected_type If a `type`, only accept final values of this type.
60 If any other callable, try to call the function on each result.
61 If the last key in the path is a `dict`, it will apply to each value inside
62 the dict instead, recursively. This does respect branching paths.
63 @param get_all If `False`, return the first matching result, otherwise all matching ones.
64 @param casesense If `False`, consider string dictionary keys as case insensitive.
66 The following are only meant to be used by YoutubeDL.prepare_outtmpl and are not part of the API
68 @param is_user_input Whether the keys are generated from user input.
69 If `True` strings get converted to `int`/`slice` if needed.
70 @param traverse_string Whether to traverse into objects as strings.
71 If `True`, any non-compatible object will first be
72 converted into a string and then traversed into.
73 The return value of that path will be a string instead,
74 not respecting any further branching.
77 @returns The result of the object traversal.
78 If successful, `get_all=True`, and the path branches at least once,
79 then a list of results is returned instead.
80 If no `default` is given and the last path branches, a `list` of results
81 is always returned. If a path ends on a `dict` that result will always be a `dict`.
83 casefold
= lambda k
: k
.casefold() if isinstance(k
, str) else k
85 if isinstance(expected_type
, type):
86 type_test
= lambda val
: val
if isinstance(val
, expected_type
) else None
88 type_test
= lambda val
: try_call(expected_type
or IDENTITY
, args
=(val
,))
90 def apply_key(key
, obj
, is_last
):
94 if obj
is None and traverse_string
:
95 if key
is ... or callable(key
) or isinstance(key
, slice):
102 elif isinstance(key
, set):
103 assert len(key
) == 1, 'Set should only be used to wrap a single item'
104 item
= next(iter(key
))
105 if isinstance(item
, type):
106 if isinstance(obj
, item
):
109 result
= try_call(item
, args
=(obj
,))
111 elif isinstance(key
, (list, tuple)):
113 result
= itertools
.chain
.from_iterable(
114 apply_path(obj
, branch
, is_last
)[0] for branch
in key
)
118 if isinstance(obj
, collections
.abc
.Mapping
):
119 result
= obj
.values()
120 elif is_iterable_like(obj
):
122 elif isinstance(obj
, re
.Match
):
123 result
= obj
.groups()
124 elif traverse_string
:
132 if isinstance(obj
, collections
.abc
.Mapping
):
133 iter_obj
= obj
.items()
134 elif is_iterable_like(obj
):
135 iter_obj
= enumerate(obj
)
136 elif isinstance(obj
, re
.Match
):
137 iter_obj
= itertools
.chain(
138 enumerate((obj
.group(), *obj
.groups())),
139 obj
.groupdict().items())
140 elif traverse_string
:
142 iter_obj
= enumerate(str(obj
))
146 result
= (v
for k
, v
in iter_obj
if try_call(key
, args
=(k
, v
)))
147 if not branching
: # string traversal
148 result
= ''.join(result
)
150 elif isinstance(key
, dict):
151 iter_obj
= ((k
, _traverse_obj(obj
, v
, False, is_last
)) for k
, v
in key
.items())
153 k
: v
if v
is not None else default
for k
, v
in iter_obj
154 if v
is not None or default
is not NO_DEFAULT
157 elif isinstance(obj
, collections
.abc
.Mapping
):
158 result
= (try_call(obj
.get
, args
=(key
,)) if casesense
or try_call(obj
.__contains
__, args
=(key
,)) else
159 next((v
for k
, v
in obj
.items() if casefold(k
) == key
), None))
161 elif isinstance(obj
, re
.Match
):
162 if isinstance(key
, int) or casesense
:
163 with contextlib
.suppress(IndexError):
164 result
= obj
.group(key
)
166 elif isinstance(key
, str):
167 result
= next((v
for k
, v
in obj
.groupdict().items() if casefold(k
) == key
), None)
169 elif isinstance(key
, (int, slice)):
170 if is_iterable_like(obj
, collections
.abc
.Sequence
):
171 branching
= isinstance(key
, slice)
172 with contextlib
.suppress(IndexError):
174 elif traverse_string
:
175 with contextlib
.suppress(IndexError):
176 result
= str(obj
)[key
]
178 return branching
, result
if branching
else (result
,)
180 def lazy_last(iterable
):
181 iterator
= iter(iterable
)
182 prev
= next(iterator
, NO_DEFAULT
)
183 if prev
is NO_DEFAULT
:
186 for item
in iterator
:
192 def apply_path(start_obj
, path
, test_type
):
197 for last
, key
in lazy_last(variadic(path
, (str, bytes, dict, set))):
198 if is_user_input
and isinstance(key
, str):
202 key
= slice(*map(int_or_none
, key
.split(':')))
203 elif int_or_none(key
) is not None:
206 if not casesense
and isinstance(key
, str):
209 if __debug__
and callable(key
):
210 # Verify function signature
211 inspect
.signature(key
).bind(None, None)
215 branching
, results
= apply_key(key
, obj
, last
)
216 has_branched |
= branching
217 new_objs
.append(results
)
219 objs
= itertools
.chain
.from_iterable(new_objs
)
221 if test_type
and not isinstance(key
, (dict, list, tuple)):
222 objs
= map(type_test
, objs
)
224 return objs
, has_branched
, isinstance(key
, dict)
226 def _traverse_obj(obj
, path
, allow_empty
, test_type
):
227 results
, has_branched
, is_dict
= apply_path(obj
, path
, test_type
)
228 results
= LazyList(item
for item
in results
if item
not in (None, {}))
229 if get_all
and has_branched
:
231 return results
.exhaust()
233 return [] if default
is NO_DEFAULT
else default
236 return results
[0] if results
else {} if allow_empty
and is_dict
else None
238 for index
, path
in enumerate(paths
, 1):
239 result
= _traverse_obj(obj
, path
, index
== len(paths
), True)
240 if result
is not None:
243 return None if default
is NO_DEFAULT
else default
246 def get_first(obj
, *paths
, **kwargs
):
247 return traverse_obj(obj
, *((..., *variadic(keys
)) for keys
in paths
), **kwargs
, get_all
=False)
250 def dict_get(d
, key_or_keys
, default
=None, skip_false_values
=True):
251 for val
in map(d
.get
, variadic(key_or_keys
)):
252 if val
is not None and (val
or not skip_false_values
):