From 40f772c3009e342a7d88f9f315d3c0b672baf768 Mon Sep 17 00:00:00 2001 From: Simon Sawicki Date: Sat, 7 Jan 2023 08:19:13 +0100 Subject: [PATCH] Rough demo --- yt_dlp/YoutubeDL.py | 41 ++++----- yt_dlp/__init__.py | 19 ++-- yt_dlp/extractor/__init__.py | 16 ++-- yt_dlp/extractor/common.py | 16 +--- yt_dlp/extractor/extractors.py | 39 ++++---- yt_dlp/globals.py | 15 ++++ yt_dlp/options.py | 10 ++- yt_dlp/plugins.py | 150 ++++++++++++++++++++++--------- yt_dlp/postprocessor/__init__.py | 28 ++++-- yt_dlp/utils.py | 4 +- 10 files changed, 213 insertions(+), 125 deletions(-) create mode 100644 yt_dlp/globals.py diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 1fb44e7f9..27eaefdc4 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -31,9 +31,15 @@ from .downloader.rtmp import rtmpdump_version from .extractor import gen_extractor_classes, get_info_extractor from .extractor.common import UnsupportedURLIE from .extractor.openload import PhantomJSwrapper +from .globals import ( + IN_CLI, + LAZY_EXTRACTORS, + plugin_ies, + plugin_overrides, + plugin_pps, +) from .minicurses import format_text from .plugins import directories as plugin_directories -from .postprocessor import _PLUGIN_CLASSES as plugin_pps from .postprocessor import ( EmbedThumbnailPP, FFmpegFixupDuplicateMoovPP, @@ -3733,15 +3739,6 @@ class YoutubeDL: if not self.params.get('verbose'): return - from . import _IN_CLI # Must be delayed import - - # These imports can be slow. So import them only as needed - from .extractor.extractors import _LAZY_LOADER - from .extractor.extractors import ( - _PLUGIN_CLASSES as plugin_ies, - _PLUGIN_OVERRIDES as plugin_ie_overrides - ) - def get_encoding(stream): ret = str(getattr(stream, 'encoding', 'missing (%s)' % type(stream).__name__)) if not supports_terminal_sequences(stream): @@ -3774,17 +3771,17 @@ class YoutubeDL: __version__, f'[{RELEASE_GIT_HEAD}]' if RELEASE_GIT_HEAD else '', '' if source == 'unknown' else f'({source})', - '' if _IN_CLI else 'API', + '' if IN_CLI.get() else 'API', delim=' ')) - if not _IN_CLI: + if not IN_CLI.get(): write_debug(f'params: {self.params}') - if not _LAZY_LOADER: - if os.environ.get('YTDLP_NO_LAZY_EXTRACTORS'): - write_debug('Lazy loading extractors is forcibly disabled') - else: - write_debug('Lazy loading extractors is disabled') + lazy_extractors = LAZY_EXTRACTORS.get() + if lazy_extractors is None: + write_debug('Lazy loading extractors is disabled') + elif not lazy_extractors: + write_debug('Lazy loading extractors is forcibly disabled') if self.params['compat_opts']: write_debug('Compatibility options: %s' % ', '.join(self.params['compat_opts'])) @@ -3818,13 +3815,13 @@ class YoutubeDL: proxy_map.update(handler.proxies) write_debug(f'Proxy map: {proxy_map}') - for plugin_type, plugins in {'Extractor': plugin_ies, 'Post-Processor': plugin_pps}.items(): - display_list = ['%s%s' % ( - klass.__name__, '' if klass.__name__ == name else f' as {name}') - for name, klass in plugins.items()] + for plugin_type, plugins in (('Extractor', plugin_ies), ('Post-Processor', plugin_pps)): + display_list = [ + klass.__name__ if klass.__name__ == name else f'{klass.__name__} as {name}' + for name, klass in plugins.get().items()] if plugin_type == 'Extractor': display_list.extend(f'{plugins[-1].IE_NAME.partition("+")[2]} ({parent.__name__})' - for parent, plugins in plugin_ie_overrides.items()) + for parent, plugins in plugin_overrides.get().items()) if not display_list: continue write_debug(f'{plugin_type} Plugins: {", ".join(sorted(display_list))}') diff --git a/yt_dlp/__init__.py b/yt_dlp/__init__.py index df1a54138..d52ed1bf3 100644 --- a/yt_dlp/__init__.py +++ b/yt_dlp/__init__.py @@ -19,8 +19,10 @@ from .cookies import SUPPORTED_BROWSERS, SUPPORTED_KEYRINGS from .downloader.external import get_external_downloader from .extractor import list_extractor_classes from .extractor.adobepass import MSO_INFO +from .globals import IN_CLI from .options import parseOpts -from .postprocessor import ( +from .plugins import load_all_plugin_types +from .postprocessor.ffmpeg import ( FFmpegExtractAudioPP, FFmpegMergerPP, FFmpegPostProcessor, @@ -28,9 +30,8 @@ from .postprocessor import ( FFmpegThumbnailsConvertorPP, FFmpegVideoConvertorPP, FFmpegVideoRemuxerPP, - MetadataFromFieldPP, - MetadataParserPP, ) +from .postprocessor.metadataparser import MetadataFromFieldPP, MetadataParserPP from .update import Updater from .utils import ( NO_DEFAULT, @@ -63,8 +64,6 @@ from .utils import ( ) from .YoutubeDL import YoutubeDL -_IN_CLI = False - def _exit(status=0, *args): for msg in args: @@ -394,6 +393,10 @@ 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 opts.playlist_items is not None: try: tuple(PlaylistEntries.parse_playlist_items(opts.playlist_items)) @@ -927,6 +930,9 @@ def _real_main(argv=None): if opts.ffmpeg_location: FFmpegPostProcessor._ffmpeg_location.set(opts.ffmpeg_location) + # load all plugins into the global lookup + load_all_plugin_types() + with YoutubeDL(ydl_opts) as ydl: pre_process = opts.update_self or opts.rm_cachedir actual_use = all_urls or opts.load_info_filename @@ -964,8 +970,7 @@ def _real_main(argv=None): def main(argv=None): - global _IN_CLI - _IN_CLI = True + IN_CLI.set(True) try: _exit(*variadic(_real_main(argv))) except DownloadError: diff --git a/yt_dlp/extractor/__init__.py b/yt_dlp/extractor/__init__.py index 6bfa4bd7b..515b7d60c 100644 --- a/yt_dlp/extractor/__init__.py +++ b/yt_dlp/extractor/__init__.py @@ -1,16 +1,16 @@ -from ..compat.compat_utils import passthrough_module +from .extractors import * +from ..globals import extractors as _extractor_classes -passthrough_module(__name__, '.extractors') -del passthrough_module +# from ..compat.compat_utils import passthrough_module +# passthrough_module(__name__, '.extractors') +# del passthrough_module def gen_extractor_classes(): """ Return a list of supported extractors. The order does matter; the first extractor matched is the one handling the URL. """ - from .extractors import _ALL_CLASSES - - return _ALL_CLASSES + return list(_extractor_classes.get().values()) def gen_extractors(): @@ -37,6 +37,4 @@ def list_extractors(age_limit=None): def get_info_extractor(ie_name): """Returns the info extractor class with the given ie_name""" - from . import extractors - - return getattr(extractors, f'{ie_name}IE') + return _extractor_classes.get()[f'{ie_name}IE'] diff --git a/yt_dlp/extractor/common.py b/yt_dlp/extractor/common.py index ef9759974..a2354983b 100644 --- a/yt_dlp/extractor/common.py +++ b/yt_dlp/extractor/common.py @@ -5,7 +5,6 @@ import hashlib import http.client import http.cookiejar import http.cookies -import inspect import itertools import json import math @@ -3724,16 +3723,8 @@ class InfoExtractor: @classmethod def __init_subclass__(cls, *, plugin_name=None, **kwargs): - if plugin_name: - mro = inspect.getmro(cls) - super_class = cls.__wrapped__ = mro[mro.index(cls) + 1] - cls.PLUGIN_NAME, cls.ie_key = plugin_name, super_class.ie_key - cls.IE_NAME = f'{super_class.IE_NAME}+{plugin_name}' - while getattr(super_class, '__wrapped__', None): - super_class = super_class.__wrapped__ - setattr(sys.modules[super_class.__module__], super_class.__name__, cls) - _PLUGIN_OVERRIDES[super_class].append(cls) - + if plugin_name is not None: + cls._plugin_name = plugin_name return super().__init_subclass__(**kwargs) @@ -3789,6 +3780,3 @@ class UnsupportedURLIE(InfoExtractor): def _real_extract(self, url): raise UnsupportedError(url) - - -_PLUGIN_OVERRIDES = collections.defaultdict(list) diff --git a/yt_dlp/extractor/extractors.py b/yt_dlp/extractor/extractors.py index baa69d242..edf2deb02 100644 --- a/yt_dlp/extractor/extractors.py +++ b/yt_dlp/extractor/extractors.py @@ -1,28 +1,25 @@ -import contextlib +import inspect import os -from ..plugins import load_plugins +from ..globals import LAZY_EXTRACTORS, extractors -# NB: Must be before other imports so that plugins can be correctly injected -_PLUGIN_CLASSES = load_plugins('extractor', 'IE') - -_LAZY_LOADER = False +_CLASS_LOOKUP = None if not os.environ.get('YTDLP_NO_LAZY_EXTRACTORS'): - with contextlib.suppress(ImportError): - from .lazy_extractors import * # noqa: F403 - from .lazy_extractors import _ALL_CLASSES - _LAZY_LOADER = True + try: + from .lazy_extractors import _CLASS_LOOKUP + LAZY_EXTRACTORS.set(True) + except ImportError: + LAZY_EXTRACTORS.set(None) -if not _LAZY_LOADER: - from ._extractors import * # noqa: F403 - _ALL_CLASSES = [ # noqa: F811 - klass - for name, klass in globals().items() - if name.endswith('IE') and name != 'GenericIE' - ] - _ALL_CLASSES.append(GenericIE) # noqa: F405 +if not _CLASS_LOOKUP: + from . import _extractors -globals().update(_PLUGIN_CLASSES) -_ALL_CLASSES[:0] = _PLUGIN_CLASSES.values() + _CLASS_LOOKUP = { + name: value + for name, value in inspect.getmembers(_extractors) + if name.endswith('IE') and name != 'GenericIE' + } + _CLASS_LOOKUP['GenericIE'] = _extractors.GenericIE -from .common import _PLUGIN_OVERRIDES # noqa: F401 +extractors.set(_CLASS_LOOKUP) +globals().update(_CLASS_LOOKUP) diff --git a/yt_dlp/globals.py b/yt_dlp/globals.py new file mode 100644 index 000000000..fb01b5f63 --- /dev/null +++ b/yt_dlp/globals.py @@ -0,0 +1,15 @@ +from collections import defaultdict +from contextvars import ContextVar + +# NAME = 'yt-dlp' + +postprocessors = ContextVar('postprocessors', default={}) +extractors = ContextVar('extractors', default={}) +IN_CLI = ContextVar('IN_CLI', default=False) +# `False`=force, `None`=disabled, `True`=enabled +LAZY_EXTRACTORS = ContextVar('LAZY_EXTRACTORS') + +plugin_dirs = ContextVar('plugin_dirs') +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 68a3aecc4..b9c1d931c 100644 --- a/yt_dlp/options.py +++ b/yt_dlp/options.py @@ -11,15 +11,15 @@ import sys from .compat import compat_expanduser from .cookies import SUPPORTED_BROWSERS, SUPPORTED_KEYRINGS from .downloader.external import list_external_downloaders -from .postprocessor import ( +from .postprocessor.ffmpeg import ( FFmpegExtractAudioPP, FFmpegMergerPP, FFmpegSubtitlesConvertorPP, FFmpegThumbnailsConvertorPP, FFmpegVideoRemuxerPP, - SponsorBlockPP, ) from .postprocessor.modify_chapters import DEFAULT_SPONSORBLOCK_CHAPTER_TITLE +from .postprocessor.sponsorblock import SponsorBlockPP from .update import detect_variant, is_non_updateable from .utils import ( OUTTMPL_TYPES, @@ -435,6 +435,12 @@ def create_parser(): '--no-colors', '--no-colours', action='store_true', dest='no_color', default=False, help='Do not emit color codes in output (Alias: --no-colours)') + general.add_option( + '--plugin-dirs', + metavar='PATH', dest='plugin_dirs', 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')) general.add_option( '--compat-options', metavar='OPTS', dest='compat_opts', default=set(), type='str', diff --git a/yt_dlp/plugins.py b/yt_dlp/plugins.py index ff5ab9d5e..43de2b5d0 100644 --- a/yt_dlp/plugins.py +++ b/yt_dlp/plugins.py @@ -1,4 +1,5 @@ import contextlib +import enum import importlib import importlib.abc import importlib.machinery @@ -12,12 +13,21 @@ import zipimport from pathlib import Path from zipfile import ZipFile +from .globals import ( + extractors, + plugin_dirs, + plugin_ies, + plugin_overrides, + plugin_pps, + postprocessors, +) + from .compat import functools # isort: split from .utils import ( get_executable_path, get_system_config_dirs, get_user_config_dirs, - orderedSet, + merge_dicts, write_string, ) @@ -25,6 +35,17 @@ PACKAGE_NAME = 'yt_dlp_plugins' COMPAT_PACKAGE_NAME = 'ytdlp_plugins' +class PluginType(enum.Enum): + POSTPROCESSORS = ('postprocessor', 'PP') + EXTRACTORS = ('extractor', 'IE') + + +_plugin_type_lookup = { + PluginType.POSTPROCESSORS: (postprocessors, plugin_pps), + PluginType.EXTRACTORS: (extractors, plugin_ies), +} + + class PluginLoader(importlib.abc.Loader): """Dummy loader for virtual namespace packages""" @@ -34,9 +55,41 @@ class PluginLoader(importlib.abc.Loader): @functools.cache def dirs_in_zip(archive): - with ZipFile(archive) as zip: - return set(itertools.chain.from_iterable( - Path(file).parents for file in zip.namelist())) + with contextlib.suppress(FileNotFoundError): + with ZipFile(archive) as zip: + return set(itertools.chain.from_iterable( + Path(file).parents for file in zip.namelist())) + return () + + +def default_plugin_paths(): + seen = set() + + def _get_unique_package_paths(*root_paths, containing_folder): + for config_dir in map(Path, root_paths): + plugin_dir = config_dir / containing_folder + # if plugin_dir in seen: + # continue + seen.add(plugin_dir) + if not plugin_dir.is_dir(): + continue + yield from plugin_dir.iterdir() + + # Load from yt-dlp config folders + yield from _get_unique_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_unique_package_paths( + get_executable_path(), + *get_user_config_dirs(''), + *get_system_config_dirs(''), + containing_folder='yt-dlp-plugins') + + # Load from PYTHONPATH folders + yield from map(Path, sys.path) class PluginFinder(importlib.abc.MetaPathFinder): @@ -53,41 +106,23 @@ class PluginFinder(importlib.abc.MetaPathFinder): 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): - plugin_dir = config_dir / containing_folder - if not plugin_dir.is_dir(): - continue - yield from plugin_dir.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 + candidate_locations = itertools.chain.from_iterable( + default_plugin_paths() if candidate is ... + else Path(candidate).iterdir() + for candidate in plugin_dirs.get((..., ))) parts = Path(*fullname.split('.')) - locations = set() + locations = dict() for path in dict.fromkeys(candidate_locations): candidate = path / parts + # print(candidate) if candidate.is_dir(): - locations.add(str(candidate)) + locations[candidate] = None elif path.name and any(path.with_suffix(suffix).is_file() for suffix in {'.zip', '.egg', '.whl'}): - with contextlib.suppress(FileNotFoundError): - if parts in dirs_in_zip(path): - locations.add(str(candidate)) - return locations + if parts in dirs_in_zip(path): + locations[candidate] = None + + return list(map(str, locations)) def find_spec(self, fullname, path=None, target=None): if fullname not in self.packages: @@ -129,7 +164,9 @@ def load_module(module, module_name, suffix): and obj.__name__ in getattr(module, '__all__', [obj.__name__]))) -def load_plugins(name, suffix): +def load_plugins(plugin_type: PluginType): + destination, plugin_destination = _plugin_type_lookup[plugin_type] + name, suffix = plugin_type.value classes = {} for finder, module_name, _ in iter_modules(name): @@ -154,17 +191,46 @@ def load_plugins(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)) + if ... in plugin_dirs.get((..., )): + 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)) + + # __init_subclass__ was removed so we manually add overrides + for name, klass in classes.items(): + plugin_name = getattr(klass, '_plugin_name', None) + if not plugin_name: + continue + + # FIXME: Most likely something wrong here + mro = inspect.getmro(klass) + super_class = klass.__wrapped__ = mro[mro.index(klass) + 1] + klass.PLUGIN_NAME, klass.ie_key = plugin_name, super_class.ie_key + klass.IE_NAME = f'{super_class.IE_NAME}+{plugin_name}' + while getattr(super_class, '__wrapped__', None): + super_class = super_class.__wrapped__ + setattr(sys.modules[super_class.__module__], super_class.__name__, klass) + plugin_overrides.get()[super_class].append(klass) + + # Add the classes into the global plugin lookup + plugin_destination.set(classes) + # We want to prepend to the main lookup + current = destination.get() + result = merge_dicts(classes, current) + destination.set(result) return classes +def load_all_plugin_types(): + for plugin_type in PluginType: + load_plugins(plugin_type) + + sys.meta_path.insert(0, PluginFinder(f'{PACKAGE_NAME}.extractor', f'{PACKAGE_NAME}.postprocessor')) -__all__ = ['directories', 'load_plugins', 'PACKAGE_NAME', 'COMPAT_PACKAGE_NAME'] +__all__ = ['directories', 'load_plugins', 'load_all_plugin_types', 'PACKAGE_NAME', 'COMPAT_PACKAGE_NAME'] diff --git a/yt_dlp/postprocessor/__init__.py b/yt_dlp/postprocessor/__init__.py index bfe9df733..0d76ba733 100644 --- a/yt_dlp/postprocessor/__init__.py +++ b/yt_dlp/postprocessor/__init__.py @@ -33,15 +33,31 @@ from .movefilesafterdownload import MoveFilesAfterDownloadPP from .sponskrub import SponSkrubPP from .sponsorblock import SponsorBlockPP from .xattrpp import XAttrMetadataPP -from ..plugins import load_plugins +from ..globals import plugin_pps, postprocessors +from ..plugins import PACKAGE_NAME +from ..utils import deprecation_warning -_PLUGIN_CLASSES = load_plugins('postprocessor', 'PP') + +def __getattr__(name): + lookup = plugin_pps.get() + if name in lookup: + deprecation_warning( + f'Importing a plugin Post-Processor from {__name__} is deprecated. ' + f'Please import {PACKAGE_NAME}.postprocessor.{name} instead.') + return lookup[name] + + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') def get_postprocessor(key): - return globals()[key + 'PP'] + return postprocessors.get()[key + 'PP'] + +_default_pps = { + name: value + for name, value in globals().items() + if name.endswith('PP') or name in ('PostProcessor', 'FFmpegPostProcessor') +} +postprocessors.set(_default_pps) -globals().update(_PLUGIN_CLASSES) -__all__ = [name for name in globals().keys() if name.endswith('PP')] -__all__.extend(('PostProcessor', 'FFmpegPostProcessor')) +__all__ = list(_default_pps.values()) diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 15e1f97cb..5fbafac99 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -56,6 +56,7 @@ from .compat import ( compat_shlex_quote, ) from .dependencies import brotli, certifi, websockets, xattr +from .globals import IN_CLI from .socks import ProxyType, sockssocket @@ -2053,8 +2054,7 @@ def write_string(s, out=None, encoding=None): def deprecation_warning(msg, *, printer=None, stacklevel=0, **kwargs): - from . import _IN_CLI - if _IN_CLI: + if IN_CLI.get(): if msg in deprecation_warning._cache: return deprecation_warning._cache.add(msg)