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