yt_dlp/plugins.py

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