yt_dlp/plugins.py

   1 import contextlib
   2 import functools
   3 import importlib
   4 import importlib.abc
   5 import importlib.machinery
   6 import importlib.util
   7 import inspect
   8 import itertools
   9 import os
  10 import pkgutil
  11 import sys
  12 import traceback
  13 import zipimport
  14 from pathlib import Path
  15 from zipfile import ZipFile
  16
  17 from .utils import (
  18     Config,
  19     get_executable_path,
  20     get_system_config_dirs,
  21     get_user_config_dirs,
  22     orderedSet,
  23     write_string,
  24 )
  25
  26 PACKAGE_NAME = 'yt_dlp_plugins'
  27 COMPAT_PACKAGE_NAME = 'ytdlp_plugins'
  28
  29
  30 class PluginLoader(importlib.abc.Loader):
  31     """Dummy loader for virtual namespace packages"""
  32
  33     def exec_module(self, module):
  34         return None
  35
  36
  37 @functools.cache
  38 def dirs_in_zip(archive):
  39     try:
  40         with ZipFile(archive) as zip_:
  41             return set(itertools.chain.from_iterable(
  42                 Path(file).parents for file in zip_.namelist()))
  43     except FileNotFoundError:
  44         pass
  45     except Exception as e:
  46         write_string(f'WARNING: Could not read zip file {archive}: {e}\n')
  47     return set()
  48
  49
  50 class PluginFinder(importlib.abc.MetaPathFinder):
  51     """
  52     This class provides one or multiple namespace packages.
  53     It searches in sys.path and yt-dlp config folders for
  54     the existing subdirectories from which the modules can be imported
  55     """
  56
  57     def __init__(self, *packages):
  58         self._zip_content_cache = {}
  59         self.packages = set(itertools.chain.from_iterable(
  60             itertools.accumulate(name.split('.'), lambda a, b: '.'.join((a, b)))
  61             for name in packages))
  62
  63     def search_locations(self, fullname):
  64         candidate_locations = []
  65
  66         def _get_package_paths(*root_paths, containing_folder='plugins'):
  67             for config_dir in orderedSet(map(Path, root_paths), lazy=True):
  68                 with contextlib.suppress(OSError):
  69                     yield from (config_dir / containing_folder).iterdir()
  70
  71         # Load from yt-dlp config folders
  72         candidate_locations.extend(_get_package_paths(
  73             *get_user_config_dirs('yt-dlp'),
  74             *get_system_config_dirs('yt-dlp'),
  75             containing_folder='plugins'))
  76
  77         # Load from yt-dlp-plugins folders
  78         candidate_locations.extend(_get_package_paths(
  79             get_executable_path(),
  80             *get_user_config_dirs(''),
  81             *get_system_config_dirs(''),
  82             containing_folder='yt-dlp-plugins'))
  83
  84         candidate_locations.extend(map(Path, sys.path))  # PYTHONPATH
  85         with contextlib.suppress(ValueError):  # Added when running __main__.py directly
  86             candidate_locations.remove(Path(__file__).parent)
  87
  88         # TODO(coletdjnz): remove when plugin globals system is implemented
  89         if Config._plugin_dirs:
  90             candidate_locations.extend(_get_package_paths(
  91                 *Config._plugin_dirs,
  92                 containing_folder=''))
  93
  94         parts = Path(*fullname.split('.'))
  95         for path in orderedSet(candidate_locations, lazy=True):
  96             candidate = path / parts
  97             try:
  98                 if candidate.is_dir():
  99                     yield candidate
 100                 elif path.suffix in ('.zip', '.egg', '.whl') and path.is_file():
 101                     if parts in dirs_in_zip(path):
 102                         yield candidate
 103             except PermissionError as e:
 104                 write_string(f'Permission error while accessing modules in "{e.filename}"\n')
 105
 106     def find_spec(self, fullname, path=None, target=None):
 107         if fullname not in self.packages:
 108             return None
 109
 110         search_locations = list(map(str, self.search_locations(fullname)))
 111         if not search_locations:
 112             return None
 113
 114         spec = importlib.machinery.ModuleSpec(fullname, PluginLoader(), is_package=True)
 115         spec.submodule_search_locations = search_locations
 116         return spec
 117
 118     def invalidate_caches(self):
 119         dirs_in_zip.cache_clear()
 120         for package in self.packages:
 121             if package in sys.modules:
 122                 del sys.modules[package]
 123
 124
 125 def directories():
 126     spec = importlib.util.find_spec(PACKAGE_NAME)
 127     return spec.submodule_search_locations if spec else []
 128
 129
 130 def iter_modules(subpackage):
 131     fullname = f'{PACKAGE_NAME}.{subpackage}'
 132     with contextlib.suppress(ModuleNotFoundError):
 133         pkg = importlib.import_module(fullname)
 134         yield from pkgutil.iter_modules(path=pkg.__path__, prefix=f'{fullname}.')
 135
 136
 137 def load_module(module, module_name, suffix):
 138     return inspect.getmembers(module, lambda obj: (
 139         inspect.isclass(obj)
 140         and obj.__name__.endswith(suffix)
 141         and obj.__module__.startswith(module_name)
 142         and not obj.__name__.startswith('_')
 143         and obj.__name__ in getattr(module, '__all__', [obj.__name__])))
 144
 145
 146 def load_plugins(name, suffix):
 147     classes = {}
 148     if os.environ.get('YTDLP_NO_PLUGINS'):
 149         return classes
 150
 151     for finder, module_name, _ in iter_modules(name):
 152         if any(x.startswith('_') for x in module_name.split('.')):
 153             continue
 154         try:
 155             if sys.version_info < (3, 10) and isinstance(finder, zipimport.zipimporter):
 156                 # zipimporter.load_module() is deprecated in 3.10 and removed in 3.12
 157                 # The exec_module branch below is the replacement for >= 3.10
 158                 # See: https://docs.python.org/3/library/zipimport.html#zipimport.zipimporter.exec_module
 159                 module = finder.load_module(module_name)
 160             else:
 161                 spec = finder.find_spec(module_name)
 162                 module = importlib.util.module_from_spec(spec)
 163                 sys.modules[module_name] = module
 164                 spec.loader.exec_module(module)
 165         except Exception:
 166             write_string(f'Error while importing module {module_name!r}\n{traceback.format_exc(limit=-1)}')
 167             continue
 168         classes.update(load_module(module, module_name, suffix))
 169
 170     # Compat: old plugin system using __init__.py
 171     # Note: plugins imported this way do not show up in directories()
 172     # nor are considered part of the yt_dlp_plugins namespace package
 173     with contextlib.suppress(FileNotFoundError):
 174         spec = importlib.util.spec_from_file_location(
 175             name, Path(get_executable_path(), COMPAT_PACKAGE_NAME, name, '__init__.py'))
 176         plugins = importlib.util.module_from_spec(spec)
 177         sys.modules[spec.name] = plugins
 178         spec.loader.exec_module(plugins)
 179         classes.update(load_module(plugins, spec.name, suffix))
 180
 181     return classes
 182
 183
 184 sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.extractor', f'{PACKAGE_NAME}.postprocessor'))
 185
 186 __all__ = ['COMPAT_PACKAGE_NAME', 'PACKAGE_NAME', 'directories', 'load_plugins']