Merge branch 'yt-dlp:master' into pr/live-sections

yt_dlp/utils/_utils.py

@@ -9,6 +9,7 @@ import datetime as dt
 import email.header
 import email.utils
 import errno
+import functools
 import hashlib
 import hmac
 import html.entities
@@ -44,7 +45,6 @@ import xml.etree.ElementTree
 
 from . import traversal
 
-from ..compat import functools  # isort: split
 from ..compat import (
     compat_etree_fromstring,
     compat_expanduser,
@@ -664,31 +664,51 @@ def sanitize_filename(s, restricted=False, is_id=NO_DEFAULT):
     return result
 
 
+def _sanitize_path_parts(parts):
+    sanitized_parts = []
+    for part in parts:
+        if not part or part == '.':
+            continue
+        elif part == '..':
+            if sanitized_parts and sanitized_parts[-1] != '..':
+                sanitized_parts.pop()
+            sanitized_parts.append('..')
+            continue
+        # Replace invalid segments with `#`
+        # - trailing dots and spaces (`asdf...` => `asdf..#`)
+        # - invalid chars (`<>` => `##`)
+        sanitized_part = re.sub(r'[/<>:"\|\\?\*]|[\s.]$', '#', part)
+        sanitized_parts.append(sanitized_part)
+
+    return sanitized_parts
+
+
 def sanitize_path(s, force=False):
     """Sanitizes and normalizes path on Windows"""
-    # XXX: this handles drive relative paths (c:sth) incorrectly
-    if sys.platform == 'win32':
-        force = False
-        drive_or_unc, _ = os.path.splitdrive(s)
-    elif force:
-        drive_or_unc = ''
-    else:
-        return s
+    if sys.platform != 'win32':
+        if not force:
+            return s
+        root = '/' if s.startswith('/') else ''
+        return root + '/'.join(_sanitize_path_parts(s.split('/')))
 
-    norm_path = os.path.normpath(remove_start(s, drive_or_unc)).split(os.path.sep)
-    if drive_or_unc:
-        norm_path.pop(0)
-    sanitized_path = [
-        path_part if path_part in ['.', '..'] else re.sub(r'(?:[/<>:"\|\\?\*]|[\s.]$)', '#', path_part)
-        for path_part in norm_path]
-    if drive_or_unc:
-        sanitized_path.insert(0, drive_or_unc + os.path.sep)
-    elif force and s and s[0] == os.path.sep:
-        sanitized_path.insert(0, os.path.sep)
-    # TODO: Fix behavioral differences <3.12
-    # The workaround using `normpath` only superficially passes tests
-    # Ref: https://github.com/python/cpython/pull/100351
-    return os.path.normpath(os.path.join(*sanitized_path))
+    normed = s.replace('/', '\\')
+
+    if normed.startswith('\\\\'):
+        # UNC path (`\\SERVER\SHARE`) or device path (`\\.`, `\\?`)
+        parts = normed.split('\\')
+        root = '\\'.join(parts[:4]) + '\\'
+        parts = parts[4:]
+    elif normed[1:2] == ':':
+        # absolute path or drive relative path
+        offset = 3 if normed[2:3] == '\\' else 2
+        root = normed[:offset]
+        parts = normed[offset:].split('\\')
+    else:
+        # relative/drive root relative path
+        root = '\\' if normed[:1] == '\\' else ''
+        parts = normed.split('\\')
+
+    return root + '\\'.join(_sanitize_path_parts(parts))
 
 
 def sanitize_url(url, *, scheme='http'):
@@ -804,14 +824,18 @@ class Popen(subprocess.Popen):
         _startupinfo = None
 
     @staticmethod
-    def _fix_pyinstaller_ld_path(env):
-        """Restore LD_LIBRARY_PATH when using PyInstaller
-            Ref: https://github.com/pyinstaller/pyinstaller/blob/develop/doc/runtime-information.rst#ld_library_path--libpath-considerations
-                 https://github.com/yt-dlp/yt-dlp/issues/4573
-        """
+    def _fix_pyinstaller_issues(env):
         if not hasattr(sys, '_MEIPASS'):
             return
 
+        # Force spawning independent subprocesses for exes bundled with PyInstaller>=6.10
+        # Ref: https://pyinstaller.org/en/v6.10.0/CHANGES.html#incompatible-changes
+        #      https://github.com/yt-dlp/yt-dlp/issues/11259
+        env['PYINSTALLER_RESET_ENVIRONMENT'] = '1'
+
+        # Restore LD_LIBRARY_PATH when using PyInstaller
+        # Ref: https://pyinstaller.org/en/v6.10.0/runtime-information.html#ld-library-path-libpath-considerations
+        #      https://github.com/yt-dlp/yt-dlp/issues/4573
         def _fix(key):
             orig = env.get(f'{key}_ORIG')
             if orig is None:
@@ -825,7 +849,7 @@ class Popen(subprocess.Popen):
     def __init__(self, args, *remaining, env=None, text=False, shell=False, **kwargs):
         if env is None:
             env = os.environ.copy()
-        self._fix_pyinstaller_ld_path(env)
+        self._fix_pyinstaller_issues(env)
 
         self.__text_mode = kwargs.get('encoding') or kwargs.get('errors') or text or kwargs.get('universal_newlines')
         if text is True:
@@ -1964,11 +1988,30 @@ def urljoin(base, path):
     return urllib.parse.urljoin(base, path)
 
 
-def int_or_none(v, scale=1, default=None, get_attr=None, invscale=1):
+def partial_application(func):
+    sig = inspect.signature(func)
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        try:
+            sig.bind(*args, **kwargs)
+        except TypeError:
+            return functools.partial(func, *args, **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+    return wrapped
+
+
+@partial_application
+def int_or_none(v, scale=1, default=None, get_attr=None, invscale=1, base=None):
     if get_attr and v is not None:
         v = getattr(v, get_attr, None)
+    if invscale == 1 and scale < 1:
+        invscale = int(1 / scale)
+        scale = 1
     try:
-        return int(v) * invscale // scale
+        return (int(v) if base is None else int(v, base=base)) * invscale // scale
     except (ValueError, TypeError, OverflowError):
         return default
 
@@ -1986,9 +2029,13 @@ def str_to_int(int_str):
         return int_or_none(int_str)
 
 
+@partial_application
 def float_or_none(v, scale=1, invscale=1, default=None):
     if v is None:
         return default
+    if invscale == 1 and scale < 1:
+        invscale = int(1 / scale)
+        scale = 1
     try:
         return float(v) * invscale / scale
     except (ValueError, TypeError):
@@ -4855,6 +4902,10 @@ class Config:
     filename = None
     __initialized = False
 
+    # Internal only, do not use! Hack to enable --plugin-dirs
+    # TODO(coletdjnz): remove when plugin globals system is implemented
+    _plugin_dirs = None
+
     def __init__(self, parser, label=None):
         self.parser, self.label = parser, label
         self._loaded_paths, self.configs = set(), []

yt_dlp/utils/traversal.py

@@ -1,18 +1,35 @@
+from __future__ import annotations
+
+import collections
 import collections.abc
 import contextlib
+import functools
 import http.cookies
 import inspect
 import itertools
 import re
+import typing
 import xml.etree.ElementTree
 
 from ._utils import (
     IDENTITY,
     NO_DEFAULT,
+    ExtractorError,
     LazyList,
     deprecation_warning,
+    get_elements_html_by_class,
+    get_elements_html_by_attribute,
+    get_elements_by_attribute,
+    get_element_html_by_attribute,
+    get_element_by_attribute,
+    get_element_html_by_id,
+    get_element_by_id,
+    get_element_html_by_class,
+    get_elements_by_class,
+    get_element_text_and_html_by_tag,
     is_iterable_like,
     try_call,
+    url_or_none,
     variadic,
 )
 
@@ -38,7 +55,7 @@ def traverse_obj(
     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}`/`{type, type, ...}/`{func}`. If a `type`, return only
+                            like `{type}`/`{type, type, ...}`/`{func}`. If a `type`, return 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]`.
@@ -54,10 +71,11 @@ def traverse_obj(
                             Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
         - `any`-builtin:    Take the first matching object and return it, resetting branching.
         - `all`-builtin:    Take all matching objects and return them as a list, resetting branching.
+        - `filter`-builtin: Return the value if it is truthy, `None` otherwise.
 
         `tuple`, `list`, and `dict` all support nested paths and branches.
 
-    @params paths           Paths which to traverse by.
+    @params paths           Paths by which to traverse.
     @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.
@@ -247,6 +265,10 @@ def traverse_obj(
                     objs = (list(filtered_objs),)
                 continue
 
+            if key is filter:
+                objs = filter(None, objs)
+                continue
+
             if __debug__ and callable(key):
                 # Verify function signature
                 inspect.signature(key).bind(None, None)
@@ -277,13 +299,143 @@ def traverse_obj(
         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
+        is_last = index == len(paths)
+        try:
+            result = _traverse_obj(obj, path, is_last, True)
+            if result is not None:
+                return result
+        except _RequiredError as e:
+            if is_last:
+                # Reraise to get cleaner stack trace
+                raise ExtractorError(e.orig_msg, expected=e.expected) from None
 
     return None if default is NO_DEFAULT else default
 
 
+def value(value, /):
+    return lambda _: value
+
+
+def require(name, /, *, expected=False):
+    def func(value):
+        if value is None:
+            raise _RequiredError(f'Unable to extract {name}', expected=expected)
+
+        return value
+
+    return func
+
+
+class _RequiredError(ExtractorError):
+    pass
+
+
+@typing.overload
+def subs_list_to_dict(*, ext: str | None = None) -> collections.abc.Callable[[list[dict]], dict[str, list[dict]]]: ...
+
+
+@typing.overload
+def subs_list_to_dict(subs: list[dict] | None, /, *, ext: str | None = None) -> dict[str, list[dict]]: ...
+
+
+def subs_list_to_dict(subs: list[dict] | None = None, /, *, ext=None):
+    """
+    Convert subtitles from a traversal into a subtitle dict.
+    The path should have an `all` immediately before this function.
+
+    Arguments:
+    `ext`      The default value for `ext` in the subtitle dict
+
+    In the dict you can set the following additional items:
+    `id`       The subtitle id to sort the dict into
+    `quality`  The sort order for each subtitle
+    """
+    if subs is None:
+        return functools.partial(subs_list_to_dict, ext=ext)
+
+    result = collections.defaultdict(list)
+
+    for sub in subs:
+        if not url_or_none(sub.get('url')) and not sub.get('data'):
+            continue
+        sub_id = sub.pop('id', None)
+        if sub_id is None:
+            continue
+        if ext is not None and not sub.get('ext'):
+            sub['ext'] = ext
+        result[sub_id].append(sub)
+    result = dict(result)
+
+    for subs in result.values():
+        subs.sort(key=lambda x: x.pop('quality', 0) or 0)
+
+    return result
+
+
+@typing.overload
+def find_element(*, attr: str, value: str, tag: str | None = None, html=False): ...
+
+
+@typing.overload
+def find_element(*, cls: str, html=False): ...
+
+
+@typing.overload
+def find_element(*, id: str, tag: str | None = None, html=False): ...
+
+
+@typing.overload
+def find_element(*, tag: str, html=False): ...
+
+
+def find_element(*, tag=None, id=None, cls=None, attr=None, value=None, html=False):
+    # deliberately using `id=` and `cls=` for ease of readability
+    assert tag or id or cls or (attr and value), 'One of tag, id, cls or (attr AND value) is required'
+    if not tag:
+        tag = r'[\w:.-]+'
+
+    if attr and value:
+        assert not cls, 'Cannot match both attr and cls'
+        assert not id, 'Cannot match both attr and id'
+        func = get_element_html_by_attribute if html else get_element_by_attribute
+        return functools.partial(func, attr, value, tag=tag)
+
+    elif cls:
+        assert not id, 'Cannot match both cls and id'
+        assert tag is None, 'Cannot match both cls and tag'
+        func = get_element_html_by_class if html else get_elements_by_class
+        return functools.partial(func, cls)
+
+    elif id:
+        func = get_element_html_by_id if html else get_element_by_id
+        return functools.partial(func, id, tag=tag)
+
+    index = int(bool(html))
+    return lambda html: get_element_text_and_html_by_tag(tag, html)[index]
+
+
+@typing.overload
+def find_elements(*, cls: str, html=False): ...
+
+
+@typing.overload
+def find_elements(*, attr: str, value: str, tag: str | None = None, html=False): ...
+
+
+def find_elements(*, tag=None, cls=None, attr=None, value=None, html=False):
+    # deliberately using `cls=` for ease of readability
+    assert cls or (attr and value), 'One of cls or (attr AND value) is required'
+
+    if attr and value:
+        assert not cls, 'Cannot match both attr and cls'
+        func = get_elements_html_by_attribute if html else get_elements_by_attribute
+        return functools.partial(func, attr, value, tag=tag or r'[\w:.-]+')
+
+    assert not tag, 'Cannot match both cls and tag'
+    func = get_elements_html_by_class if html else get_elements_by_class
+    return functools.partial(func, cls)
+
+
 def get_first(obj, *paths, **kwargs):
     return traverse_obj(obj, *((..., *variadic(keys)) for keys in paths), **kwargs, get_all=False)