devscripts/make_lazy_extractors.py

   1 #!/usr/bin/env python3
   2
   3 # Allow direct execution
   4 import os
   5 import shutil
   6 import sys
   7
   8 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
   9
  10
  11 from inspect import getsource
  12
  13 from devscripts.utils import get_filename_args, read_file, write_file
  14
  15 NO_ATTR = object()
  16 STATIC_CLASS_PROPERTIES = [
  17     'IE_NAME', '_ENABLED', '_VALID_URL',  # Used for URL matching
  18     '_WORKING', 'IE_DESC', '_NETRC_MACHINE', 'SEARCH_KEY',  # Used for --extractor-descriptions
  19     'age_limit',  # Used for --age-limit (evaluated)
  20     '_RETURN_TYPE',  # Accessed in CLI only with instance (evaluated)
  21 ]
  22 CLASS_METHODS = [
  23     'ie_key', 'suitable', '_match_valid_url',  # Used for URL matching
  24     'working', 'get_temp_id', '_match_id',  # Accessed just before instance creation
  25     'description',  # Used for --extractor-descriptions
  26     'is_suitable',  # Used for --age-limit
  27     'supports_login', 'is_single_video',  # Accessed in CLI only with instance
  28 ]
  29 IE_TEMPLATE = '''
  30 class {name}({bases}):
  31     _module = {module!r}
  32 '''
  33 MODULE_TEMPLATE = read_file('devscripts/lazy_load_template.py')
  34
  35
  36 def main():
  37     lazy_extractors_filename = get_filename_args(default_outfile='yt_dlp/extractor/lazy_extractors.py')
  38     if os.path.exists(lazy_extractors_filename):
  39         os.remove(lazy_extractors_filename)
  40
  41     _ALL_CLASSES = get_all_ies()  # Must be before import
  42
  43     import yt_dlp.plugins
  44     from yt_dlp.extractor.common import InfoExtractor, SearchInfoExtractor
  45
  46     # Filter out plugins
  47     _ALL_CLASSES = [cls for cls in _ALL_CLASSES if not cls.__module__.startswith(f'{yt_dlp.plugins.PACKAGE_NAME}.')]
  48
  49     DummyInfoExtractor = type('InfoExtractor', (InfoExtractor,), {'IE_NAME': NO_ATTR})
  50     module_src = '\n'.join((
  51         MODULE_TEMPLATE,
  52         '    _module = None',
  53         *extra_ie_code(DummyInfoExtractor),
  54         '\nclass LazyLoadSearchExtractor(LazyLoadExtractor):\n    pass\n',
  55         *build_ies(_ALL_CLASSES, (InfoExtractor, SearchInfoExtractor), DummyInfoExtractor),
  56     ))
  57
  58     write_file(lazy_extractors_filename, f'{module_src}\n')
  59
  60
  61 def get_all_ies():
  62     PLUGINS_DIRNAME = 'ytdlp_plugins'
  63     BLOCKED_DIRNAME = f'{PLUGINS_DIRNAME}_blocked'
  64     if os.path.exists(PLUGINS_DIRNAME):
  65         # os.rename cannot be used, e.g. in Docker. See https://github.com/yt-dlp/yt-dlp/pull/4958
  66         shutil.move(PLUGINS_DIRNAME, BLOCKED_DIRNAME)
  67     try:
  68         from yt_dlp.extractor.extractors import _ALL_CLASSES
  69     finally:
  70         if os.path.exists(BLOCKED_DIRNAME):
  71             shutil.move(BLOCKED_DIRNAME, PLUGINS_DIRNAME)
  72     return _ALL_CLASSES
  73
  74
  75 def extra_ie_code(ie, base=None):
  76     for var in STATIC_CLASS_PROPERTIES:
  77         val = getattr(ie, var)
  78         if val != (getattr(base, var) if base else NO_ATTR):
  79             yield f'    {var} = {val!r}'
  80     yield ''
  81
  82     for name in CLASS_METHODS:
  83         f = getattr(ie, name)
  84         if not base or f.__func__ != getattr(base, name).__func__:
  85             yield getsource(f)
  86
  87
  88 def build_ies(ies, bases, attr_base):
  89     names = []
  90     for ie in sort_ies(ies, bases):
  91         yield build_lazy_ie(ie, ie.__name__, attr_base)
  92         if ie in ies:
  93             names.append(ie.__name__)
  94
  95     yield f'\n_ALL_CLASSES = [{", ".join(names)}]'
  96
  97
  98 def sort_ies(ies, ignored_bases):
  99     """find the correct sorting and add the required base classes so that subclasses can be correctly created"""
 100     classes, returned_classes = ies[:-1], set()
 101     assert ies[-1].__name__ == 'GenericIE', 'Last IE must be GenericIE'
 102     while classes:
 103         for c in classes[:]:
 104             bases = set(c.__bases__) - {object, *ignored_bases}
 105             restart = False
 106             for b in sorted(bases, key=lambda x: x.__name__):
 107                 if b not in classes and b not in returned_classes:
 108                     assert b.__name__ != 'GenericIE', 'Cannot inherit from GenericIE'
 109                     classes.insert(0, b)
 110                     restart = True
 111             if restart:
 112                 break
 113             if bases <= returned_classes:
 114                 yield c
 115                 returned_classes.add(c)
 116                 classes.remove(c)
 117                 break
 118     yield ies[-1]
 119
 120
 121 def build_lazy_ie(ie, name, attr_base):
 122     bases = ', '.join({
 123         'InfoExtractor': 'LazyLoadExtractor',
 124         'SearchInfoExtractor': 'LazyLoadSearchExtractor',
 125     }.get(base.__name__, base.__name__) for base in ie.__bases__)
 126
 127     s = IE_TEMPLATE.format(name=name, module=ie.__module__, bases=bases)
 128     return s + '\n'.join(extra_ie_code(ie, attr_base))
 129
 130
 131 if __name__ == '__main__':
 132     main()