255 lines
10 KiB
Python
255 lines
10 KiB
Python
|
import collections.abc
|
||
|
|
import contextlib
|
||
|
|
import inspect
|
||
|
|
import itertools
|
||
|
|
import re
|
||
|
|
|
||
|
|
from ._utils import (
|
||
|
|
IDENTITY,
|
||
|
|
NO_DEFAULT,
|
||
|
|
LazyList,
|
||
|
|
int_or_none,
|
||
|
|
is_iterable_like,
|
||
|
|
try_call,
|
||
|
|
variadic,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def traverse_obj(
|
||
|
|
obj, *paths, default=NO_DEFAULT, expected_type=None, get_all=True,
|
||
|
|
casesense=True, is_user_input=False, traverse_string=False):
|
||
|
|
"""
|
||
|
|
Safely traverse nested `dict`s and `Iterable`s
|
||
|
|
|
||
|
|
>>> obj = [{}, {"key": "value"}]
|
||
|
|
>>> traverse_obj(obj, (1, "key"))
|
||
|
|
"value"
|
||
|
|
|
||
|
|
Each of the provided `paths` is tested and the first producing a valid result will be returned.
|
||
|
|
The next path will also be tested if the path branched but no results could be found.
|
||
|
|
Supported values for traversal are `Mapping`, `Iterable` and `re.Match`.
|
||
|
|
Unhelpful values (`{}`, `None`) are treated as the absence of a value and discarded.
|
||
|
|
|
||
|
|
The paths will be wrapped in `variadic`, so that `'key'` is conveniently the same as `('key', )`.
|
||
|
|
|
||
|
|
The keys in the path can be one of:
|
||
|
|
- `None`: Return the current object.
|
||
|
|
- `set`: Requires the only item in the set to be a type or function,
|
||
|
|
like `{type}`/`{func}`. If a `type`, returns only values
|
||
|
|
of this type. If a function, returns `func(obj)`.
|
||
|
|
- `str`/`int`: Return `obj[key]`. For `re.Match`, return `obj.group(key)`.
|
||
|
|
- `slice`: Branch out and return all values in `obj[key]`.
|
||
|
|
- `Ellipsis`: Branch out and return a list of all values.
|
||
|
|
- `tuple`/`list`: Branch out and return a list of all matching values.
|
||
|
|
Read as: `[traverse_obj(obj, branch) for branch in branches]`.
|
||
|
|
- `function`: Branch out and return values filtered by the function.
|
||
|
|
Read as: `[value for key, value in obj if function(key, value)]`.
|
||
|
|
For `Iterable`s, `key` is the index of the value.
|
||
|
|
For `re.Match`es, `key` is the group number (0 = full match)
|
||
|
|
as well as additionally any group names, if given.
|
||
|
|
- `dict` Transform the current object and return a matching dict.
|
||
|
|
Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
|
||
|
|
|
||
|
|
`tuple`, `list`, and `dict` all support nested paths and branches.
|
||
|
|
|
||
|
|
@params paths Paths which to traverse by.
|
||
|
|
@param default Value to return if the paths do not match.
|
||
|
|
If the last key in the path is a `dict`, it will apply to each value inside
|
||
|
|
the dict instead, depth first. Try to avoid if using nested `dict` keys.
|
||
|
|
@param expected_type If a `type`, only accept final values of this type.
|
||
|
|
If any other callable, try to call the function on each result.
|
||
|
|
If the last key in the path is a `dict`, it will apply to each value inside
|
||
|
|
the dict instead, recursively. This does respect branching paths.
|
||
|
|
@param get_all If `False`, return the first matching result, otherwise all matching ones.
|
||
|
|
@param casesense If `False`, consider string dictionary keys as case insensitive.
|
||
|
|
|
||
|
|
The following are only meant to be used by YoutubeDL.prepare_outtmpl and are not part of the API
|
||
|
|
|
||
|
|
@param is_user_input Whether the keys are generated from user input.
|
||
|
|
If `True` strings get converted to `int`/`slice` if needed.
|
||
|
|
@param traverse_string Whether to traverse into objects as strings.
|
||
|
|
If `True`, any non-compatible object will first be
|
||
|
|
converted into a string and then traversed into.
|
||
|
|
The return value of that path will be a string instead,
|
||
|
|
not respecting any further branching.
|
||
|
|
|
||
|
|
|
||
|
|
@returns The result of the object traversal.
|
||
|
|
If successful, `get_all=True`, and the path branches at least once,
|
||
|
|
then a list of results is returned instead.
|
||
|
|
If no `default` is given and the last path branches, a `list` of results
|
||
|
|
is always returned. If a path ends on a `dict` that result will always be a `dict`.
|
||
|
|
"""
|
||
|
|
casefold = lambda k: k.casefold() if isinstance(k, str) else k
|
||
|
|
|
||
|
|
if isinstance(expected_type, type):
|
||
|
|
type_test = lambda val: val if isinstance(val, expected_type) else None
|
||
|
|
else:
|
||
|
|
type_test = lambda val: try_call(expected_type or IDENTITY, args=(val,))
|
||
|
|
|
||
|
|
def apply_key(key, obj, is_last):
|
||
|
|
branching = False
|
||
|
|
result = None
|
||
|
|
|
||
|
|
if obj is None and traverse_string:
|
||
|
|
if key is ... or callable(key) or isinstance(key, slice):
|
||
|
|
branching = True
|
||
|
|
result = ()
|
||
|
|
|
||
|
|
elif key is None:
|
||
|
|
result = obj
|
||
|
|
|
||
|
|
elif isinstance(key, set):
|
||
|
|
assert len(key) == 1, 'Set should only be used to wrap a single item'
|
||
|
|
item = next(iter(key))
|
||
|
|
if isinstance(item, type):
|
||
|
|
if isinstance(obj, item):
|
||
|
|
result = obj
|
||
|
|
else:
|
||
|
|
result = try_call(item, args=(obj,))
|
||
|
|
|
||
|
|
elif isinstance(key, (list, tuple)):
|
||
|
|
branching = True
|
||
|
|
result = itertools.chain.from_iterable(
|
||
|
|
apply_path(obj, branch, is_last)[0] for branch in key)
|
||
|
|
|
||
|
|
elif key is ...:
|
||
|
|
branching = True
|
||
|
|
if isinstance(obj, collections.abc.Mapping):
|
||
|
|
result = obj.values()
|
||
|
|
elif is_iterable_like(obj):
|
||
|
|
result = obj
|
||
|
|
elif isinstance(obj, re.Match):
|
||
|
|
result = obj.groups()
|
||
|
|
elif traverse_string:
|
||
|
|
branching = False
|
||
|
|
result = str(obj)
|
||
|
|
else:
|
||
|
|
result = ()
|
||
|
|
|
||
|
|
elif callable(key):
|
||
|
|
branching = True
|
||
|
|
if isinstance(obj, collections.abc.Mapping):
|
||
|
|
iter_obj = obj.items()
|
||
|
|
elif is_iterable_like(obj):
|
||
|
|
iter_obj = enumerate(obj)
|
||
|
|
elif isinstance(obj, re.Match):
|
||
|
|
iter_obj = itertools.chain(
|
||
|
|
enumerate((obj.group(), *obj.groups())),
|
||
|
|
obj.groupdict().items())
|
||
|
|
elif traverse_string:
|
||
|
|
branching = False
|
||
|
|
iter_obj = enumerate(str(obj))
|
||
|
|
else:
|
||
|
|
iter_obj = ()
|
||
|
|
|
||
|
|
result = (v for k, v in iter_obj if try_call(key, args=(k, v)))
|
||
|
|
if not branching: # string traversal
|
||
|
|
result = ''.join(result)
|
||
|
|
|
||
|
|
elif isinstance(key, dict):
|
||
|
|
iter_obj = ((k, _traverse_obj(obj, v, False, is_last)) for k, v in key.items())
|
||
|
|
result = {
|
||
|
|
k: v if v is not None else default for k, v in iter_obj
|
||
|
|
if v is not None or default is not NO_DEFAULT
|
||
|
|
} or None
|
||
|
|
|
||
|
|
elif isinstance(obj, collections.abc.Mapping):
|
||
|
|
result = (try_call(obj.get, args=(key,)) if casesense or try_call(obj.__contains__, args=(key,)) else
|
||
|
|
next((v for k, v in obj.items() if casefold(k) == key), None))
|
||
|
|
|
||
|
|
elif isinstance(obj, re.Match):
|
||
|
|
if isinstance(key, int) or casesense:
|
||
|
|
with contextlib.suppress(IndexError):
|
||
|
|
result = obj.group(key)
|
||
|
|
|
||
|
|
elif isinstance(key, str):
|
||
|
|
result = next((v for k, v in obj.groupdict().items() if casefold(k) == key), None)
|
||
|
|
|
||
|
|
elif isinstance(key, (int, slice)):
|
||
|
|
if is_iterable_like(obj, collections.abc.Sequence):
|
||
|
|
branching = isinstance(key, slice)
|
||
|
|
with contextlib.suppress(IndexError):
|
||
|
|
result = obj[key]
|
||
|
|
elif traverse_string:
|
||
|
|
with contextlib.suppress(IndexError):
|
||
|
|
result = str(obj)[key]
|
||
|
|
|
||
|
|
return branching, result if branching else (result,)
|
||
|
|
|
||
|
|
def lazy_last(iterable):
|
||
|
|
iterator = iter(iterable)
|
||
|
|
prev = next(iterator, NO_DEFAULT)
|
||
|
|
if prev is NO_DEFAULT:
|
||
|
|
return
|
||
|
|
|
||
|
|
for item in iterator:
|
||
|
|
yield False, prev
|
||
|
|
prev = item
|
||
|
|
|
||
|
|
yield True, prev
|
||
|
|
|
||
|
|
def apply_path(start_obj, path, test_type):
|
||
|
|
objs = (start_obj,)
|
||
|
|
has_branched = False
|
||
|
|
|
||
|
|
key = None
|
||
|
|
for last, key in lazy_last(variadic(path, (str, bytes, dict, set))):
|
||
|
|
if is_user_input and isinstance(key, str):
|
||
|
|
if key == ':':
|
||
|
|
key = ...
|
||
|
|
elif ':' in key:
|
||
|
|
key = slice(*map(int_or_none, key.split(':')))
|
||
|
|
elif int_or_none(key) is not None:
|
||
|
|
key = int(key)
|
||
|
|
|
||
|
|
if not casesense and isinstance(key, str):
|
||
|
|
key = key.casefold()
|
||
|
|
|
||
|
|
if __debug__ and callable(key):
|
||
|
|
# Verify function signature
|
||
|
|
inspect.signature(key).bind(None, None)
|
||
|
|
|
||
|
|
new_objs = []
|
||
|
|
for obj in objs:
|
||
|
|
branching, results = apply_key(key, obj, last)
|
||
|
|
has_branched |= branching
|
||
|
|
new_objs.append(results)
|
||
|
|
|
||
|
|
objs = itertools.chain.from_iterable(new_objs)
|
||
|
|
|
||
|
|
if test_type and not isinstance(key, (dict, list, tuple)):
|
||
|
|
objs = map(type_test, objs)
|
||
|
|
|
||
|
|
return objs, has_branched, isinstance(key, dict)
|
||
|
|
|
||
|
|
def _traverse_obj(obj, path, allow_empty, test_type):
|
||
|
|
results, has_branched, is_dict = apply_path(obj, path, test_type)
|
||
|
|
results = LazyList(item for item in results if item not in (None, {}))
|
||
|
|
if get_all and has_branched:
|
||
|
|
if results:
|
||
|
|
return results.exhaust()
|
||
|
|
if allow_empty:
|
||
|
|
return [] if default is NO_DEFAULT else default
|
||
|
|
return None
|
||
|
|
|
||
|
|
return results[0] if results else {} if allow_empty and is_dict else None
|
||
|
|
|
||
|
|
for index, path in enumerate(paths, 1):
|
||
|
|
result = _traverse_obj(obj, path, index == len(paths), True)
|
||
|
|
if result is not None:
|
||
|
|
return result
|
||
|
|
|
||
|
|
return None if default is NO_DEFAULT else default
|
||
|
|
|
||
|
|
|
||
|
|
def get_first(obj, *paths, **kwargs):
|
||
|
|
return traverse_obj(obj, *((..., *variadic(keys)) for keys in paths), **kwargs, get_all=False)
|
||
|
|
|
||
|
|
|
||
|
|
def dict_get(d, key_or_keys, default=None, skip_false_values=True):
|
||
|
|
for val in map(d.get, variadic(key_or_keys)):
|
||
|
|
if val is not None and (val or not skip_false_values):
|
||
|
|
return val
|
||
|
|
return default
|