diff --git a/test/test_YoutubeDL.py b/test/test_YoutubeDL.py index ece57d2ce1..cbb60f6b15 100644 --- a/test/test_YoutubeDL.py +++ b/test/test_YoutubeDL.py @@ -6,7 +6,7 @@ import sys import unittest from unittest.mock import patch -from yt_dlp._globals import ALL_PLUGINS_LOADED +from yt_dlp._globals import all_plugins_loaded sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -1400,9 +1400,9 @@ class TestYoutubeDL(unittest.TestCase): def test_load_plugins_compat(self): # Should try to reload plugins if they haven't already been loaded - ALL_PLUGINS_LOADED.set(False) + all_plugins_loaded.set(False) FakeYDL().close() - assert ALL_PLUGINS_LOADED.get() + assert all_plugins_loaded.get() if __name__ == '__main__': diff --git a/test/test_plugins.py b/test/test_plugins.py index 0956ef7497..a57a2b3c27 100644 --- a/test/test_plugins.py +++ b/test/test_plugins.py @@ -11,8 +11,23 @@ TEST_DATA_DIR = Path(os.path.dirname(os.path.abspath(__file__)), 'testdata') sys.path.append(str(TEST_DATA_DIR)) importlib.invalidate_caches() -from yt_dlp.plugins import PACKAGE_NAME, PluginType, directories, load_plugins, load_all_plugin_types -from yt_dlp._globals import extractors, postprocessors, plugin_dirs, plugin_ies, plugin_pps, ALL_PLUGINS_LOADED +from yt_dlp.plugins import PACKAGE_NAME, PluginSpec, directories, load_plugins, load_all_plugins, register_plugin_spec +from yt_dlp._globals import extractors, postprocessors, plugin_dirs, plugin_ies, plugin_pps, all_plugins_loaded, plugin_specs + + +EXTRACTOR_PLUGIN_SPEC = PluginSpec( + module_name='extractor', + suffix='IE', + destination=extractors, + plugin_destination=plugin_ies, +) + +POSTPROCESSOR_PLUGIN_SPEC = PluginSpec( + module_name='postprocessor', + suffix='PP', + destination=postprocessors, + plugin_destination=plugin_pps, +) class TestPlugins(unittest.TestCase): @@ -23,7 +38,8 @@ class TestPlugins(unittest.TestCase): plugin_ies.set({}) plugin_pps.set({}) plugin_dirs.set((...,)) - ALL_PLUGINS_LOADED.set(False) + plugin_specs.set({}) + all_plugins_loaded.set(False) importlib.invalidate_caches() # Clearing override plugins is probably difficult for module_name in tuple(sys.modules): @@ -35,7 +51,7 @@ class TestPlugins(unittest.TestCase): self.assertIn(self.TEST_PLUGIN_DIR, map(Path, directories())) def test_extractor_classes(self): - plugins_ie = load_plugins(PluginType.EXTRACTORS) + plugins_ie = load_plugins(EXTRACTOR_PLUGIN_SPEC) self.assertIn(f'{PACKAGE_NAME}.extractor.normal', sys.modules.keys()) self.assertIn('NormalPluginIE', plugins_ie.keys()) @@ -64,7 +80,7 @@ class TestPlugins(unittest.TestCase): self.assertNotIn('_UnderscoreOverrideGenericIE', plugin_ies.get()) def test_postprocessor_classes(self): - plugins_pp = load_plugins(PluginType.POSTPROCESSORS) + plugins_pp = load_plugins(POSTPROCESSOR_PLUGIN_SPEC) self.assertIn('NormalPluginPP', plugins_pp.keys()) self.assertIn(f'{PACKAGE_NAME}.postprocessor.normal', sys.modules.keys()) self.assertIn('NormalPluginPP', plugin_pps.get()) @@ -80,10 +96,10 @@ class TestPlugins(unittest.TestCase): package = importlib.import_module(f'{PACKAGE_NAME}.{plugin_type}') self.assertIn(zip_path / PACKAGE_NAME / plugin_type, map(Path, package.__path__)) - plugins_ie = load_plugins(PluginType.EXTRACTORS) + plugins_ie = load_plugins(EXTRACTOR_PLUGIN_SPEC) self.assertIn('ZippedPluginIE', plugins_ie.keys()) - plugins_pp = load_plugins(PluginType.POSTPROCESSORS) + plugins_pp = load_plugins(POSTPROCESSOR_PLUGIN_SPEC) self.assertIn('ZippedPluginPP', plugins_pp.keys()) finally: @@ -93,11 +109,8 @@ class TestPlugins(unittest.TestCase): def test_reloading_plugins(self): reload_plugins_path = TEST_DATA_DIR / 'reload_plugins' - - for plugin_type in ('extractor', 'postprocessor'): - importlib.import_module(f'{PACKAGE_NAME}.{plugin_type}') - load_plugins(PluginType.EXTRACTORS) - load_plugins(PluginType.POSTPROCESSORS) + load_plugins(EXTRACTOR_PLUGIN_SPEC) + load_plugins(POSTPROCESSOR_PLUGIN_SPEC) # Remove default folder and add reload_plugin path sys.path.remove(str(TEST_DATA_DIR)) @@ -108,7 +121,7 @@ class TestPlugins(unittest.TestCase): package = importlib.import_module(f'{PACKAGE_NAME}.{plugin_type}') self.assertIn(reload_plugins_path / PACKAGE_NAME / plugin_type, map(Path, package.__path__)) - plugins_ie = load_plugins(PluginType.EXTRACTORS) + plugins_ie = load_plugins(EXTRACTOR_PLUGIN_SPEC) self.assertIn('NormalPluginIE', plugins_ie.keys()) self.assertTrue( plugins_ie['NormalPluginIE'].REPLACED, @@ -117,7 +130,7 @@ class TestPlugins(unittest.TestCase): extractors.get()['NormalPluginIE'].REPLACED, msg='Reloading has not replaced original extractor plugin globally') - plugins_pp = load_plugins(PluginType.POSTPROCESSORS) + plugins_pp = load_plugins(POSTPROCESSOR_PLUGIN_SPEC) self.assertIn('NormalPluginPP', plugins_pp.keys()) self.assertTrue(plugins_pp['NormalPluginPP'].REPLACED, msg='Reloading has not replaced original postprocessor plugin') @@ -131,7 +144,7 @@ class TestPlugins(unittest.TestCase): importlib.invalidate_caches() def test_extractor_override_plugin(self): - load_plugins(PluginType.EXTRACTORS) + load_plugins(EXTRACTOR_PLUGIN_SPEC) from yt_dlp.extractor.generic import GenericIE @@ -141,25 +154,29 @@ class TestPlugins(unittest.TestCase): self.assertEqual(GenericIE.IE_NAME, 'generic+override+underscore-override') importlib.invalidate_caches() # test that loading a second time doesn't wrap a second time - load_plugins(PluginType.EXTRACTORS) + load_plugins(EXTRACTOR_PLUGIN_SPEC) from yt_dlp.extractor.generic import GenericIE self.assertEqual(GenericIE.IE_NAME, 'generic+override+underscore-override') def test_load_all_plugin_types(self): + # no plugin specs registered + load_all_plugins() + self.assertNotIn(f'{PACKAGE_NAME}.extractor.normal', sys.modules.keys()) self.assertNotIn(f'{PACKAGE_NAME}.postprocessor.normal', sys.modules.keys()) - load_all_plugin_types() - self.assertTrue(yt_dlp._globals.ALL_PLUGINS_LOADED.get()) + register_plugin_spec(EXTRACTOR_PLUGIN_SPEC) + register_plugin_spec(POSTPROCESSOR_PLUGIN_SPEC) + load_all_plugins() + self.assertTrue(yt_dlp._globals.all_plugins_loaded.get()) 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'))) - load_all_plugin_types() - self.assertTrue(yt_dlp._globals.ALL_PLUGINS_LOADED.get()) + load_plugins(EXTRACTOR_PLUGIN_SPEC) self.assertIn(f'{PACKAGE_NAME}.extractor.package', sys.modules.keys()) self.assertIn('PackagePluginIE', plugin_ies.get()) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 2928b43876..94b62138ac 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -38,7 +38,8 @@ from ._globals import ( LAZY_EXTRACTORS, plugin_ies, plugin_overrides, - plugin_pps, ALL_PLUGINS_LOADED, + plugin_pps, + all_plugins_loaded, ) from .minicurses import format_text from .networking import HEADRequest, Request, RequestDirector @@ -51,7 +52,7 @@ from .networking.exceptions import ( network_exceptions, ) from .networking.impersonate import ImpersonateRequestHandler -from .plugins import directories as plugin_directories, load_all_plugin_types +from .plugins import directories as plugin_directories, load_all_plugins from .postprocessor import ( EmbedThumbnailPP, FFmpegFixupDuplicateMoovPP, @@ -646,8 +647,8 @@ class YoutubeDL: self.__header_cookies = [] # compat for API: load plugins if they have not already - if not ALL_PLUGINS_LOADED.get(): - load_all_plugin_types() + if not all_plugins_loaded.get(): + load_all_plugins() stdout = sys.stderr if self.params.get('logtostderr') else sys.stdout self._out_files = Namespace( diff --git a/yt_dlp/__init__.py b/yt_dlp/__init__.py index 360599b9c1..174cffcc4d 100644 --- a/yt_dlp/__init__.py +++ b/yt_dlp/__init__.py @@ -22,7 +22,7 @@ from .extractor.adobepass import MSO_INFO from .networking.impersonate import ImpersonateTarget from ._globals import IN_CLI, plugin_dirs from .options import parseOpts -from .plugins import load_all_plugin_types +from .plugins import load_all_plugins from .postprocessor import ( FFmpegExtractAudioPP, FFmpegMergerPP, @@ -987,7 +987,7 @@ def _real_main(argv=None): # load all plugins into the global lookup plugin_dirs.set(opts.plugin_dirs) - load_all_plugin_types() + load_all_plugins() with YoutubeDL(ydl_opts) as ydl: pre_process = opts.update_self or opts.rm_cachedir diff --git a/yt_dlp/_globals.py b/yt_dlp/_globals.py index b7b462e5cd..4e5a6f3e85 100644 --- a/yt_dlp/_globals.py +++ b/yt_dlp/_globals.py @@ -9,8 +9,13 @@ IN_CLI = ContextVar('IN_CLI', default=False) # `False`=force, `None`=disabled, `True`=enabled LAZY_EXTRACTORS = ContextVar('LAZY_EXTRACTORS', default=False) + +# Plugins + +plugin_specs = ContextVar('plugin_specs', default={}) + # Whether plugins have been loaded once -ALL_PLUGINS_LOADED = ContextVar('PLUGINS_LOADED', default=False) +all_plugins_loaded = ContextVar('all_plugins_loaded', default=False) # `...`=search default plugin dirs plugin_dirs = ContextVar('plugin_dirs', default=(..., )) diff --git a/yt_dlp/extractor/__init__.py b/yt_dlp/extractor/__init__.py index 75e6592159..31dc488b65 100644 --- a/yt_dlp/extractor/__init__.py +++ b/yt_dlp/extractor/__init__.py @@ -1,9 +1,18 @@ from .._globals import extractors as _extractors_context +from .._globals import plugin_ies as _plugin_ies_context from ..compat.compat_utils import passthrough_module +from ..plugins import PluginSpec, register_plugin_spec passthrough_module(__name__, '.extractors') del passthrough_module +register_plugin_spec(PluginSpec( + module_name='extractor', + suffix='IE', + destination=_extractors_context, + plugin_destination=_plugin_ies_context, +)) + def gen_extractor_classes(): """ Return a list of supported extractors. diff --git a/yt_dlp/plugins.py b/yt_dlp/plugins.py index f09787e9d6..3ba2aadf2e 100644 --- a/yt_dlp/plugins.py +++ b/yt_dlp/plugins.py @@ -1,6 +1,5 @@ import contextlib import dataclasses -import enum import importlib import importlib.abc import importlib.machinery @@ -17,12 +16,9 @@ from pathlib import Path from zipfile import ZipFile from ._globals import ( - extractors, plugin_dirs, - plugin_ies, - plugin_pps, - postprocessors, - ALL_PLUGINS_LOADED, + all_plugins_loaded, + plugin_specs, ) from .compat import functools # isort: split @@ -40,9 +36,12 @@ COMPAT_PACKAGE_NAME = 'ytdlp_plugins' _BASE_PACKAGE_PATH = Path(__file__).parent -class PluginType(enum.Enum): - POSTPROCESSORS = ('postprocessor', 'PP') - EXTRACTORS = ('extractor', 'IE') +@dataclasses.dataclass +class PluginSpec: + module_name: str + suffix: str + destination: ContextVar + plugin_destination: ContextVar class PluginLoader(importlib.abc.Loader): @@ -68,7 +67,7 @@ def dirs_in_zip(archive): return () -def default_plugin_paths(): +def external_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 @@ -115,7 +114,7 @@ class PluginFinder(importlib.abc.MetaPathFinder): def search_locations(self, fullname): candidate_locations = itertools.chain.from_iterable( - default_plugin_paths() if candidate is ... else Path(candidate).iterdir() + external_plugin_paths() if candidate is ... else Path(candidate).iterdir() for candidate in plugin_dirs.get() ) @@ -174,27 +173,8 @@ def get_regular_classes(module, module_name, suffix): )) -@dataclasses.dataclass -class _PluginTypeConfig: - destination: ContextVar - plugin_destination: ContextVar - - -_plugin_type_lookup = { - PluginType.POSTPROCESSORS: _PluginTypeConfig( - destination=postprocessors, - plugin_destination=plugin_pps, - ), - PluginType.EXTRACTORS: _PluginTypeConfig( - destination=extractors, - plugin_destination=plugin_ies, - ), -} - - -def load_plugins(plugin_type: PluginType): - plugin_config = _plugin_type_lookup[plugin_type] - name, suffix = plugin_type.value +def load_plugins(plugin_spec: PluginSpec): + name, suffix = plugin_spec.module_name, plugin_spec.suffix regular_classes = {} if os.environ.get('YTDLP_NO_PLUGINS'): return regular_classes @@ -235,25 +215,36 @@ def load_plugins(plugin_type: PluginType): regular_classes.update(get_regular_classes(plugins, spec.name, suffix)) # Add the classes into the global plugin lookup for that type - plugin_config.plugin_destination.set(regular_classes) + plugin_spec.plugin_destination.set(regular_classes) # We want to prepend to the main lookup for that type - plugin_config.destination.set(merge_dicts(regular_classes, plugin_config.destination.get())) + plugin_spec.destination.set(merge_dicts(regular_classes, plugin_spec.destination.get())) return regular_classes -def load_all_plugin_types(): - for plugin_type in PluginType: - load_plugins(plugin_type) - ALL_PLUGINS_LOADED.set(True) +def load_all_plugins(): + for plugin_spec in plugin_specs.get().values(): + load_plugins(plugin_spec) + all_plugins_loaded.set(True) -sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.extractor', f'{PACKAGE_NAME}.postprocessor')) +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.get(): + plugin_specs.get()[plugin_spec.module_name] = plugin_spec + 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) + __all__ = [ 'directories', 'load_plugins', - 'load_all_plugin_types', + 'load_all_plugins', + 'register_plugin_spec', + 'get_plugin_spec', 'PACKAGE_NAME', 'COMPAT_PACKAGE_NAME', ] diff --git a/yt_dlp/postprocessor/__init__.py b/yt_dlp/postprocessor/__init__.py index 866fe3d1c4..36ac6c7ea1 100644 --- a/yt_dlp/postprocessor/__init__.py +++ b/yt_dlp/postprocessor/__init__.py @@ -34,7 +34,7 @@ from .sponskrub import SponSkrubPP from .sponsorblock import SponsorBlockPP from .xattrpp import XAttrMetadataPP from .._globals import plugin_pps, postprocessors -from ..plugins import PACKAGE_NAME +from ..plugins import PACKAGE_NAME, register_plugin_spec, PluginSpec from ..utils import deprecation_warning @@ -53,6 +53,13 @@ def get_postprocessor(key): return postprocessors.get()[key + 'PP'] +register_plugin_spec(PluginSpec( + module_name='postprocessor', + suffix='PP', + destination=postprocessors, + plugin_destination=plugin_pps, +)) + _default_pps = { name: value for name, value in globals().items()