]>
jfr.im git - yt-dlp.git/blob - yt_dlp/utils/traversal.py
6 import xml
.etree
.ElementTree
20 obj
, *paths
, default
=NO_DEFAULT
, expected_type
=None, get_all
=True,
21 casesense
=True, is_user_input
=NO_DEFAULT
, traverse_string
=False):
23 Safely traverse nested `dict`s and `Iterable`s
25 >>> obj = [{}, {"key": "value"}]
26 >>> traverse_obj(obj, (1, "key"))
29 Each of the provided `paths` is tested and the first producing a valid result will be returned.
30 The next path will also be tested if the path branched but no results could be found.
31 Supported values for traversal are `Mapping`, `Iterable` and `re.Match`.
32 Unhelpful values (`{}`, `None`) are treated as the absence of a value and discarded.
34 The paths will be wrapped in `variadic`, so that `'key'` is conveniently the same as `('key', )`.
36 The keys in the path can be one of:
37 - `None`: Return the current object.
38 - `set`: Requires the only item in the set to be a type or function,
39 like `{type}`/`{func}`. If a `type`, returns only values
40 of this type. If a function, returns `func(obj)`.
41 - `str`/`int`: Return `obj[key]`. For `re.Match`, return `obj.group(key)`.
42 - `slice`: Branch out and return all values in `obj[key]`.
43 - `Ellipsis`: Branch out and return a list of all values.
44 - `tuple`/`list`: Branch out and return a list of all matching values.
45 Read as: `[traverse_obj(obj, branch) for branch in branches]`.
46 - `function`: Branch out and return values filtered by the function.
47 Read as: `[value for key, value in obj if function(key, value)]`.
48 For `Iterable`s, `key` is the index of the value.
49 For `re.Match`es, `key` is the group number (0 = full match)
50 as well as additionally any group names, if given.
51 - `dict` Transform the current object and return a matching dict.
52 Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
54 `tuple`, `list`, and `dict` all support nested paths and branches.
56 @params paths Paths which to traverse by.
57 @param default Value to return if the paths do not match.
58 If the last key in the path is a `dict`, it will apply to each value inside
59 the dict instead, depth first. Try to avoid if using nested `dict` keys.
60 @param expected_type If a `type`, only accept final values of this type.
61 If any other callable, try to call the function on each result.
62 If the last key in the path is a `dict`, it will apply to each value inside
63 the dict instead, recursively. This does respect branching paths.
64 @param get_all If `False`, return the first matching result, otherwise all matching ones.
65 @param casesense If `False`, consider string dictionary keys as case insensitive.
67 `traverse_string` is only meant to be used by YoutubeDL.prepare_outtmpl and is not part of the API
69 @param traverse_string Whether to traverse into objects as strings.
70 If `True`, any non-compatible object will first be
71 converted into a string and then traversed into.
72 The return value of that path will be a string instead,
73 not respecting any further branching.
76 @returns The result of the object traversal.
77 If successful, `get_all=True`, and the path branches at least once,
78 then a list of results is returned instead.
79 If no `default` is given and the last path branches, a `list` of results
80 is always returned. If a path ends on a `dict` that result will always be a `dict`.
82 if is_user_input
is not NO_DEFAULT
:
83 deprecation_warning('The is_user_input parameter is deprecated and no longer works')
85 casefold
= lambda k
: k
.casefold() if isinstance(k
, str) else k
87 if isinstance(expected_type
, type):
88 type_test
= lambda val
: val
if isinstance(val
, expected_type
) else None
90 type_test
= lambda val
: try_call(expected_type
or IDENTITY
, args
=(val
,))
92 def apply_key(key
, obj
, is_last
):
96 if obj
is None and traverse_string
:
97 if key
is ... or callable(key
) or isinstance(key
, slice):
104 elif isinstance(key
, set):
105 assert len(key
) == 1, 'Set should only be used to wrap a single item'
106 item
= next(iter(key
))
107 if isinstance(item
, type):
108 if isinstance(obj
, item
):
111 result
= try_call(item
, args
=(obj
,))
113 elif isinstance(key
, (list, tuple)):
115 result
= itertools
.chain
.from_iterable(
116 apply_path(obj
, branch
, is_last
)[0] for branch
in key
)
120 if isinstance(obj
, collections
.abc
.Mapping
):
121 result
= obj
.values()
122 elif is_iterable_like(obj
) or isinstance(obj
, xml
.etree
.ElementTree
.Element
):
124 elif isinstance(obj
, re
.Match
):
125 result
= obj
.groups()
126 elif traverse_string
:
134 if isinstance(obj
, collections
.abc
.Mapping
):
135 iter_obj
= obj
.items()
136 elif is_iterable_like(obj
) or isinstance(obj
, xml
.etree
.ElementTree
.Element
):
137 iter_obj
= enumerate(obj
)
138 elif isinstance(obj
, re
.Match
):
139 iter_obj
= itertools
.chain(
140 enumerate((obj
.group(), *obj
.groups())),
141 obj
.groupdict().items())
142 elif traverse_string
:
144 iter_obj
= enumerate(str(obj
))
148 result
= (v
for k
, v
in iter_obj
if try_call(key
, args
=(k
, v
)))
149 if not branching
: # string traversal
150 result
= ''.join(result
)
152 elif isinstance(key
, dict):
153 iter_obj
= ((k
, _traverse_obj(obj
, v
, False, is_last
)) for k
, v
in key
.items())
155 k
: v
if v
is not None else default
for k
, v
in iter_obj
156 if v
is not None or default
is not NO_DEFAULT
159 elif isinstance(obj
, collections
.abc
.Mapping
):
160 result
= (try_call(obj
.get
, args
=(key
,)) if casesense
or try_call(obj
.__contains
__, args
=(key
,)) else
161 next((v
for k
, v
in obj
.items() if casefold(k
) == key
), None))
163 elif isinstance(obj
, re
.Match
):
164 if isinstance(key
, int) or casesense
:
165 with contextlib
.suppress(IndexError):
166 result
= obj
.group(key
)
168 elif isinstance(key
, str):
169 result
= next((v
for k
, v
in obj
.groupdict().items() if casefold(k
) == key
), None)
171 elif isinstance(key
, (int, slice)):
172 if is_iterable_like(obj
, (collections
.abc
.Sequence
, xml
.etree
.ElementTree
.Element
)):
173 branching
= isinstance(key
, slice)
174 with contextlib
.suppress(IndexError):
176 elif traverse_string
:
177 with contextlib
.suppress(IndexError):
178 result
= str(obj
)[key
]
180 elif isinstance(obj
, xml
.etree
.ElementTree
.Element
) and isinstance(key
, str):
181 xpath
, _
, special
= key
.rpartition('/')
182 if not special
.startswith('@') and special
!= 'text()':
186 # Allow abbreviations of relative paths, absolute paths error
187 if xpath
.startswith('/'):
189 elif xpath
and not xpath
.startswith('./'):
192 def apply_specials(element
):
196 return element
.attrib
197 if special
.startswith('@'):
198 return try_call(element
.attrib
.get
, args
=(special
[1:],))
199 if special
== 'text()':
201 assert False, f
'apply_specials is missing case for {special!r}'
204 result
= list(map(apply_specials
, obj
.iterfind(xpath
)))
206 result
= apply_specials(obj
)
208 return branching
, result
if branching
else (result
,)
210 def lazy_last(iterable
):
211 iterator
= iter(iterable
)
212 prev
= next(iterator
, NO_DEFAULT
)
213 if prev
is NO_DEFAULT
:
216 for item
in iterator
:
222 def apply_path(start_obj
, path
, test_type
):
227 for last
, key
in lazy_last(variadic(path
, (str, bytes, dict, set))):
228 if not casesense
and isinstance(key
, str):
231 if key
in (any
, all
):
233 filtered_objs
= (obj
for obj
in objs
if obj
not in (None, {}))
235 objs
= (next(filtered_objs
, None),)
237 objs
= (list(filtered_objs
),)
240 if __debug__
and callable(key
):
241 # Verify function signature
242 inspect
.signature(key
).bind(None, None)
246 branching
, results
= apply_key(key
, obj
, last
)
247 has_branched |
= branching
248 new_objs
.append(results
)
250 objs
= itertools
.chain
.from_iterable(new_objs
)
252 if test_type
and not isinstance(key
, (dict, list, tuple)):
253 objs
= map(type_test
, objs
)
255 return objs
, has_branched
, isinstance(key
, dict)
257 def _traverse_obj(obj
, path
, allow_empty
, test_type
):
258 results
, has_branched
, is_dict
= apply_path(obj
, path
, test_type
)
259 results
= LazyList(item
for item
in results
if item
not in (None, {}))
260 if get_all
and has_branched
:
262 return results
.exhaust()
264 return [] if default
is NO_DEFAULT
else default
267 return results
[0] if results
else {} if allow_empty
and is_dict
else None
269 for index
, path
in enumerate(paths
, 1):
270 result
= _traverse_obj(obj
, path
, index
== len(paths
), True)
271 if result
is not None:
274 return None if default
is NO_DEFAULT
else default
277 def get_first(obj
, *paths
, **kwargs
):
278 return traverse_obj(obj
, *((..., *variadic(keys
)) for keys
in paths
), **kwargs
, get_all
=False)
281 def dict_get(d
, key_or_keys
, default
=None, skip_false_values
=True):
282 for val
in map(d
.get
, variadic(key_or_keys
)):
283 if val
is not None and (val
or not skip_false_values
):