yt_dlp/utils/traversal.py

   1 import collections.abc
   2 import contextlib
   3 import inspect
   4 import itertools
   5 import re
   6
   7 from ._utils import (
   8     IDENTITY,
   9     NO_DEFAULT,
  10     LazyList,
  11     int_or_none,
  12     is_iterable_like,
  13     try_call,
  14     variadic,
  15 )
  16
  17
  18 def traverse_obj(
  19         obj, *paths, default=NO_DEFAULT, expected_type=None, get_all=True,
  20         casesense=True, is_user_input=False, traverse_string=False):
  21     """
  22     Safely traverse nested `dict`s and `Iterable`s
  23
  24     >>> obj = [{}, {"key": "value"}]
  25     >>> traverse_obj(obj, (1, "key"))
  26     "value"
  27
  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.
  32
  33     The paths will be wrapped in `variadic`, so that `'key'` is conveniently the same as `('key', )`.
  34
  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()}`.
  52
  53         `tuple`, `list`, and `dict` all support nested paths and branches.
  54
  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.
  65
  66     The following are only meant to be used by YoutubeDL.prepare_outtmpl and are not part of the API
  67
  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.
  75
  76
  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`.
  82     """
  83     casefold = lambda k: k.casefold() if isinstance(k, str) else k
  84
  85     if isinstance(expected_type, type):
  86         type_test = lambda val: val if isinstance(val, expected_type) else None
  87     else:
  88         type_test = lambda val: try_call(expected_type or IDENTITY, args=(val,))
  89
  90     def apply_key(key, obj, is_last):
  91         branching = False
  92         result = None
  93
  94         if obj is None and traverse_string:
  95             if key is ... or callable(key) or isinstance(key, slice):
  96                 branching = True
  97                 result = ()
  98
  99         elif key is None:
 100             result = obj
 101
 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):
 107                     result = obj
 108             else:
 109                 result = try_call(item, args=(obj,))
 110
 111         elif isinstance(key, (list, tuple)):
 112             branching = True
 113             result = itertools.chain.from_iterable(
 114                 apply_path(obj, branch, is_last)[0] for branch in key)
 115
 116         elif key is ...:
 117             branching = True
 118             if isinstance(obj, collections.abc.Mapping):
 119                 result = obj.values()
 120             elif is_iterable_like(obj):
 121                 result = obj
 122             elif isinstance(obj, re.Match):
 123                 result = obj.groups()
 124             elif traverse_string:
 125                 branching = False
 126                 result = str(obj)
 127             else:
 128                 result = ()
 129
 130         elif callable(key):
 131             branching = True
 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:
 141                 branching = False
 142                 iter_obj = enumerate(str(obj))
 143             else:
 144                 iter_obj = ()
 145
 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)
 149
 150         elif isinstance(key, dict):
 151             iter_obj = ((k, _traverse_obj(obj, v, False, is_last)) for k, v in key.items())
 152             result = {
 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
 155             } or None
 156
 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))
 160
 161         elif isinstance(obj, re.Match):
 162             if isinstance(key, int) or casesense:
 163                 with contextlib.suppress(IndexError):
 164                     result = obj.group(key)
 165
 166             elif isinstance(key, str):
 167                 result = next((v for k, v in obj.groupdict().items() if casefold(k) == key), None)
 168
 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):
 173                     result = obj[key]
 174             elif traverse_string:
 175                 with contextlib.suppress(IndexError):
 176                     result = str(obj)[key]
 177
 178         return branching, result if branching else (result,)
 179
 180     def lazy_last(iterable):
 181         iterator = iter(iterable)
 182         prev = next(iterator, NO_DEFAULT)
 183         if prev is NO_DEFAULT:
 184             return
 185
 186         for item in iterator:
 187             yield False, prev
 188             prev = item
 189
 190         yield True, prev
 191
 192     def apply_path(start_obj, path, test_type):
 193         objs = (start_obj,)
 194         has_branched = False
 195
 196         key = None
 197         for last, key in lazy_last(variadic(path, (str, bytes, dict, set))):
 198             if is_user_input and isinstance(key, str):
 199                 if key == ':':
 200                     key = ...
 201                 elif ':' in key:
 202                     key = slice(*map(int_or_none, key.split(':')))
 203                 elif int_or_none(key) is not None:
 204                     key = int(key)
 205
 206             if not casesense and isinstance(key, str):
 207                 key = key.casefold()
 208
 209             if __debug__ and callable(key):
 210                 # Verify function signature
 211                 inspect.signature(key).bind(None, None)
 212
 213             new_objs = []
 214             for obj in objs:
 215                 branching, results = apply_key(key, obj, last)
 216                 has_branched |= branching
 217                 new_objs.append(results)
 218
 219             objs = itertools.chain.from_iterable(new_objs)
 220
 221         if test_type and not isinstance(key, (dict, list, tuple)):
 222             objs = map(type_test, objs)
 223
 224         return objs, has_branched, isinstance(key, dict)
 225
 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:
 230             if results:
 231                 return results.exhaust()
 232             if allow_empty:
 233                 return [] if default is NO_DEFAULT else default
 234             return None
 235
 236         return results[0] if results else {} if allow_empty and is_dict else None
 237
 238     for index, path in enumerate(paths, 1):
 239         result = _traverse_obj(obj, path, index == len(paths), True)
 240         if result is not None:
 241             return result
 242
 243     return None if default is NO_DEFAULT else default
 244
 245
 246 def get_first(obj, *paths, **kwargs):
 247     return traverse_obj(obj, *((..., *variadic(keys)) for keys in paths), **kwargs, get_all=False)
 248
 249
 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):
 253             return val
 254     return default