From 109c019e8a5e84889387a9f2f2d07aade870b295 Mon Sep 17 00:00:00 2001 From: coletdjnz Date: Sun, 20 Oct 2024 12:48:02 +1300 Subject: [PATCH] Add public functions to add custom external plugin paths --- test/test_plugins.py | 24 ++++++++++++++++++++--- yt_dlp/__init__.py | 16 ++++++++------- yt_dlp/_globals.py | 3 +-- yt_dlp/options.py | 2 +- yt_dlp/plugins.py | 46 +++++++++++++++++++++++++++++++------------- 5 files changed, 65 insertions(+), 26 deletions(-) diff --git a/test/test_plugins.py b/test/test_plugins.py index a57a2b3c2..00a4c6506 100644 --- a/test/test_plugins.py +++ b/test/test_plugins.py @@ -5,6 +5,7 @@ import sys import unittest from pathlib import Path import yt_dlp._globals +from yt_dlp.plugins import set_plugin_dirs, add_plugin_dirs, PluginDirs sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) TEST_DATA_DIR = Path(os.path.dirname(os.path.abspath(__file__)), 'testdata') @@ -37,7 +38,7 @@ class TestPlugins(unittest.TestCase): def setUp(self): plugin_ies.set({}) plugin_pps.set({}) - plugin_dirs.set((...,)) + plugin_dirs.set((PluginDirs.DEFAULT_EXTERNAL,)) plugin_specs.set({}) all_plugins_loaded.set(False) importlib.invalidate_caches() @@ -174,8 +175,25 @@ class TestPlugins(unittest.TestCase): self.assertIn(f'{PACKAGE_NAME}.extractor.normal', sys.modules.keys()) self.assertIn(f'{PACKAGE_NAME}.postprocessor.normal', sys.modules.keys()) - def test_plugin_dirs(self): - plugin_dirs.set((..., str(TEST_DATA_DIR / 'plugin_packages'))) + def test_set_plugin_dirs(self): + + custom_plugin_dir = str(TEST_DATA_DIR / 'plugin_packages') + set_plugin_dirs(custom_plugin_dir) + + self.assertEqual(plugin_dirs.get(), (custom_plugin_dir, )) + self.assertNotIn('external', plugin_dirs.get()) + load_plugins(EXTRACTOR_PLUGIN_SPEC) + + self.assertIn(f'{PACKAGE_NAME}.extractor.package', sys.modules.keys()) + self.assertIn('PackagePluginIE', plugin_ies.get()) + + def test_add_plugin_dirs(self): + custom_plugin_dir = str(TEST_DATA_DIR / 'plugin_packages') + + self.assertEqual(plugin_dirs.get(), (PluginDirs.DEFAULT_EXTERNAL,)) + add_plugin_dirs(custom_plugin_dir) + self.assertEqual(plugin_dirs.get(), (PluginDirs.DEFAULT_EXTERNAL, custom_plugin_dir)) + load_plugins(EXTRACTOR_PLUGIN_SPEC) self.assertIn(f'{PACKAGE_NAME}.extractor.package', sys.modules.keys()) diff --git a/yt_dlp/__init__.py b/yt_dlp/__init__.py index 174cffcc4..f7948a233 100644 --- a/yt_dlp/__init__.py +++ b/yt_dlp/__init__.py @@ -20,9 +20,11 @@ from .downloader.external import get_external_downloader from .extractor import list_extractor_classes from .extractor.adobepass import MSO_INFO from .networking.impersonate import ImpersonateTarget -from ._globals import IN_CLI, plugin_dirs +from ._globals import IN_CLI as _IN_CLI from .options import parseOpts -from .plugins import load_all_plugins +from .plugins import load_all_plugins as _load_all_plugins +from .plugins import PluginDirs as _PluginDirs +from .plugins import set_plugin_dirs as _set_plugin_dirs from .postprocessor import ( FFmpegExtractAudioPP, FFmpegMergerPP, @@ -428,8 +430,8 @@ def validate_options(opts): # Other options opts.plugin_dirs = opts.plugin_dirs or [] - if 'no-default' not in opts.plugin_dirs: - opts.plugin_dirs.append(...) + if 'no-external' not in opts.plugin_dirs: + opts.plugin_dirs.append(_PluginDirs.DEFAULT_EXTERNAL) if opts.playlist_items is not None: try: @@ -986,8 +988,8 @@ def _real_main(argv=None): FFmpegPostProcessor._ffmpeg_location.set(opts.ffmpeg_location) # load all plugins into the global lookup - plugin_dirs.set(opts.plugin_dirs) - load_all_plugins() + _set_plugin_dirs(*opts.plugin_dirs) + _load_all_plugins() with YoutubeDL(ydl_opts) as ydl: pre_process = opts.update_self or opts.rm_cachedir @@ -1088,7 +1090,7 @@ def _real_main(argv=None): def main(argv=None): - IN_CLI.set(True) + _IN_CLI.set(True) try: _exit(*variadic(_real_main(argv))) except (CookieLoadError, DownloadError): diff --git a/yt_dlp/_globals.py b/yt_dlp/_globals.py index 4e5a6f3e8..4eae01fb6 100644 --- a/yt_dlp/_globals.py +++ b/yt_dlp/_globals.py @@ -17,8 +17,7 @@ plugin_specs = ContextVar('plugin_specs', default={}) # Whether plugins have been loaded once all_plugins_loaded = ContextVar('all_plugins_loaded', default=False) -# `...`=search default plugin dirs -plugin_dirs = ContextVar('plugin_dirs', default=(..., )) +plugin_dirs = ContextVar('plugin_dirs', default=('external', )) plugin_ies = ContextVar('plugin_ies', default={}) plugin_overrides = ContextVar('plugin_overrides', default=defaultdict(list)) plugin_pps = ContextVar('plugin_pps', default={}) diff --git a/yt_dlp/options.py b/yt_dlp/options.py index f6f39c9d8..43c2760ca 100644 --- a/yt_dlp/options.py +++ b/yt_dlp/options.py @@ -472,7 +472,7 @@ def create_parser(): action='append', help=( 'Directory to search for plugins. Can be used multiple times to add multiple directories. ' - 'Add "no-default" to disable the default plugin directories' + 'Add "no-external" to disable searching default external plugin directories (outside of python environment)' ), ) general.add_option( diff --git a/yt_dlp/plugins.py b/yt_dlp/plugins.py index 3ba2aadf2..952816a62 100644 --- a/yt_dlp/plugins.py +++ b/yt_dlp/plugins.py @@ -1,5 +1,6 @@ import contextlib import dataclasses +import enum import importlib import importlib.abc import importlib.machinery @@ -36,6 +37,26 @@ COMPAT_PACKAGE_NAME = 'ytdlp_plugins' _BASE_PACKAGE_PATH = Path(__file__).parent +# Public APIs +# Anything else is NOT public and no backwards compatibility is guaranteed +__all__ = [ + 'directories', + 'load_plugins', + 'load_all_plugins', + 'register_plugin_spec', + 'add_plugin_dirs', + 'set_plugin_dirs', + 'PluginDirs', + 'get_plugin_spec', + 'PACKAGE_NAME', + 'COMPAT_PACKAGE_NAME', +] + + +class PluginDirs(enum.Enum): + DEFAULT_EXTERNAL = 'external' # The default external plugin directories + + @dataclasses.dataclass class PluginSpec: module_name: str @@ -114,7 +135,7 @@ class PluginFinder(importlib.abc.MetaPathFinder): def search_locations(self, fullname): candidate_locations = itertools.chain.from_iterable( - external_plugin_paths() if candidate is ... else Path(candidate).iterdir() + external_plugin_paths() if candidate == PluginDirs.DEFAULT_EXTERNAL else Path(candidate).iterdir() for candidate in plugin_dirs.get() ) @@ -203,7 +224,7 @@ def load_plugins(plugin_spec: PluginSpec): # 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 - if ... in plugin_dirs.get((...,)): + if PluginDirs.DEFAULT_EXTERNAL in plugin_dirs.get(): with contextlib.suppress(FileNotFoundError): spec = importlib.util.spec_from_file_location( name, @@ -235,16 +256,15 @@ def register_plugin_spec(plugin_spec: PluginSpec): sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.{plugin_spec.module_name}')) -def get_plugin_spec(module_name): - return plugin_specs.get().get(module_name) +def add_plugin_dirs(*paths): + """Add external plugin dirs to the existing ones""" + plugin_dirs.set((*plugin_dirs.get(), *paths)) -__all__ = [ - 'directories', - 'load_plugins', - 'load_all_plugins', - 'register_plugin_spec', - 'get_plugin_spec', - 'PACKAGE_NAME', - 'COMPAT_PACKAGE_NAME', -] +def set_plugin_dirs(*paths): + """Set external plugin dirs, overriding the default ones""" + plugin_dirs.set(tuple(paths)) + + +def get_plugin_spec(module_name): + return plugin_specs.get().get(module_name)