diff options
Diffstat (limited to 'yt_dlp/plugins.py')
| -rw-r--r-- | yt_dlp/plugins.py | 186 |
1 files changed, 127 insertions, 59 deletions
diff --git a/yt_dlp/plugins.py b/yt_dlp/plugins.py index 94335a9a3..941709b21 100644 --- a/yt_dlp/plugins.py +++ b/yt_dlp/plugins.py @@ -1,4 +1,5 @@ import contextlib +import dataclasses import functools import importlib import importlib.abc @@ -14,17 +15,48 @@ import zipimport from pathlib import Path from zipfile import ZipFile +from .globals import ( + Indirect, + plugin_dirs, + all_plugins_loaded, + plugin_specs, +) + from .utils import ( - Config, get_executable_path, get_system_config_dirs, get_user_config_dirs, + merge_dicts, orderedSet, write_string, ) PACKAGE_NAME = 'yt_dlp_plugins' COMPAT_PACKAGE_NAME = 'ytdlp_plugins' +_BASE_PACKAGE_PATH = Path(__file__).parent + + +# Please Note: Due to necessary changes and the complex nature involved, +# no backwards compatibility is guaranteed for the plugin system API. +# However, we will still try our best. + +__all__ = [ + 'COMPAT_PACKAGE_NAME', + 'PACKAGE_NAME', + 'PluginSpec', + 'directories', + 'load_all_plugins', + 'load_plugins', + 'register_plugin_spec', +] + + +@dataclasses.dataclass +class PluginSpec: + module_name: str + suffix: str + destination: Indirect + plugin_destination: Indirect class PluginLoader(importlib.abc.Loader): @@ -44,7 +76,42 @@ def dirs_in_zip(archive): pass except Exception as e: write_string(f'WARNING: Could not read zip file {archive}: {e}\n') - return set() + return () + + +def default_plugin_paths(): + def _get_package_paths(*root_paths, containing_folder): + for config_dir in orderedSet(map(Path, root_paths), lazy=True): + # We need to filter the base path added when running __main__.py directly + if config_dir == _BASE_PACKAGE_PATH: + continue + with contextlib.suppress(OSError): + yield from (config_dir / containing_folder).iterdir() + + # Load from yt-dlp config folders + yield from _get_package_paths( + *get_user_config_dirs('yt-dlp'), + *get_system_config_dirs('yt-dlp'), + containing_folder='plugins', + ) + + # Load from yt-dlp-plugins folders + yield from _get_package_paths( + get_executable_path(), + *get_user_config_dirs(''), + *get_system_config_dirs(''), + containing_folder='yt-dlp-plugins', + ) + + # Load from PYTHONPATH directories + yield from (path for path in map(Path, sys.path) if path != _BASE_PACKAGE_PATH) + + +def candidate_plugin_paths(candidate): + candidate_path = Path(candidate) + if not candidate_path.is_dir(): + raise ValueError(f'Invalid plugin directory: {candidate_path}') + yield from candidate_path.iterdir() class PluginFinder(importlib.abc.MetaPathFinder): @@ -56,40 +123,16 @@ class PluginFinder(importlib.abc.MetaPathFinder): def __init__(self, *packages): self._zip_content_cache = {} - self.packages = set(itertools.chain.from_iterable( - itertools.accumulate(name.split('.'), lambda a, b: '.'.join((a, b))) - for name in packages)) + self.packages = set( + itertools.chain.from_iterable( + itertools.accumulate(name.split('.'), lambda a, b: '.'.join((a, b))) + for name in packages)) def search_locations(self, fullname): - candidate_locations = [] - - def _get_package_paths(*root_paths, containing_folder='plugins'): - for config_dir in orderedSet(map(Path, root_paths), lazy=True): - with contextlib.suppress(OSError): - yield from (config_dir / containing_folder).iterdir() - - # Load from yt-dlp config folders - candidate_locations.extend(_get_package_paths( - *get_user_config_dirs('yt-dlp'), - *get_system_config_dirs('yt-dlp'), - containing_folder='plugins')) - - # Load from yt-dlp-plugins folders - candidate_locations.extend(_get_package_paths( - get_executable_path(), - *get_user_config_dirs(''), - *get_system_config_dirs(''), - containing_folder='yt-dlp-plugins')) - - candidate_locations.extend(map(Path, sys.path)) # PYTHONPATH - with contextlib.suppress(ValueError): # Added when running __main__.py directly - candidate_locations.remove(Path(__file__).parent) - - # TODO(coletdjnz): remove when plugin globals system is implemented - if Config._plugin_dirs: - candidate_locations.extend(_get_package_paths( - *Config._plugin_dirs, - containing_folder='')) + candidate_locations = itertools.chain.from_iterable( + default_plugin_paths() if candidate == 'default' else candidate_plugin_paths(candidate) + for candidate in plugin_dirs.value + ) parts = Path(*fullname.split('.')) for path in orderedSet(candidate_locations, lazy=True): @@ -109,7 +152,8 @@ class PluginFinder(importlib.abc.MetaPathFinder): search_locations = list(map(str, self.search_locations(fullname))) if not search_locations: - return None + # Prevent using built-in meta finders for searching plugins. + raise ModuleNotFoundError(fullname) spec = importlib.machinery.ModuleSpec(fullname, PluginLoader(), is_package=True) spec.submodule_search_locations = search_locations @@ -123,8 +167,10 @@ class PluginFinder(importlib.abc.MetaPathFinder): def directories(): - spec = importlib.util.find_spec(PACKAGE_NAME) - return spec.submodule_search_locations if spec else [] + with contextlib.suppress(ModuleNotFoundError): + if spec := importlib.util.find_spec(PACKAGE_NAME): + return list(spec.submodule_search_locations) + return [] def iter_modules(subpackage): @@ -134,19 +180,23 @@ def iter_modules(subpackage): yield from pkgutil.iter_modules(path=pkg.__path__, prefix=f'{fullname}.') -def load_module(module, module_name, suffix): +def get_regular_classes(module, module_name, suffix): + # Find standard public plugin classes (not overrides) return inspect.getmembers(module, lambda obj: ( inspect.isclass(obj) and obj.__name__.endswith(suffix) and obj.__module__.startswith(module_name) and not obj.__name__.startswith('_') - and obj.__name__ in getattr(module, '__all__', [obj.__name__]))) + and obj.__name__ in getattr(module, '__all__', [obj.__name__]) + and getattr(obj, 'PLUGIN_NAME', None) is None + )) -def load_plugins(name, suffix): - classes = {} - if os.environ.get('YTDLP_NO_PLUGINS'): - return classes +def load_plugins(plugin_spec: PluginSpec): + name, suffix = plugin_spec.module_name, plugin_spec.suffix + regular_classes = {} + if os.environ.get('YTDLP_NO_PLUGINS') or not plugin_dirs.value: + return regular_classes for finder, module_name, _ in iter_modules(name): if any(x.startswith('_') for x in module_name.split('.')): @@ -163,24 +213,42 @@ def load_plugins(name, suffix): sys.modules[module_name] = module spec.loader.exec_module(module) except Exception: - write_string(f'Error while importing module {module_name!r}\n{traceback.format_exc(limit=-1)}') + write_string( + f'Error while importing module {module_name!r}\n{traceback.format_exc(limit=-1)}', + ) continue - classes.update(load_module(module, module_name, suffix)) + regular_classes.update(get_regular_classes(module, module_name, suffix)) # Compat: old plugin system using __init__.py # Note: plugins imported this way do not show up in directories() # nor are considered part of the yt_dlp_plugins namespace package - with contextlib.suppress(FileNotFoundError): - spec = importlib.util.spec_from_file_location( - name, Path(get_executable_path(), COMPAT_PACKAGE_NAME, name, '__init__.py')) - plugins = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = plugins - spec.loader.exec_module(plugins) - classes.update(load_module(plugins, spec.name, suffix)) - - return classes - - -sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.extractor', f'{PACKAGE_NAME}.postprocessor')) - -__all__ = ['COMPAT_PACKAGE_NAME', 'PACKAGE_NAME', 'directories', 'load_plugins'] + if 'default' in plugin_dirs.value: + with contextlib.suppress(FileNotFoundError): + spec = importlib.util.spec_from_file_location( + name, + Path(get_executable_path(), COMPAT_PACKAGE_NAME, name, '__init__.py'), + ) + plugins = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = plugins + spec.loader.exec_module(plugins) + regular_classes.update(get_regular_classes(plugins, spec.name, suffix)) + + # Add the classes into the global plugin lookup for that type + plugin_spec.plugin_destination.value = regular_classes + # We want to prepend to the main lookup for that type + plugin_spec.destination.value = merge_dicts(regular_classes, plugin_spec.destination.value) + + return regular_classes + + +def load_all_plugins(): + for plugin_spec in plugin_specs.value.values(): + load_plugins(plugin_spec) + all_plugins_loaded.value = True + + +def register_plugin_spec(plugin_spec: PluginSpec): + # If the plugin spec for a module is already registered, it will not be added again + if plugin_spec.module_name not in plugin_specs.value: + plugin_specs.value[plugin_spec.module_name] = plugin_spec + sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.{plugin_spec.module_name}')) |