devscripts/make_lazy_extractors.py

   1 #!/usr/bin/env python3
   2 import optparse
   3 import os
   4 import sys
   5 from inspect import getsource
   6
   7 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
   8
   9
  10 NO_ATTR = object()
  11 STATIC_CLASS_PROPERTIES = ['IE_NAME', 'IE_DESC', 'SEARCH_KEY', '_WORKING', '_NETRC_MACHINE', 'age_limit']
  12 CLASS_METHODS = [
  13     'ie_key', 'working', 'description', 'suitable', '_match_valid_url', '_match_id', 'get_temp_id', 'is_suitable'
  14 ]
  15 IE_TEMPLATE = '''
  16 class {name}({bases}):
  17     _module = {module!r}
  18 '''
  19 with open('devscripts/lazy_load_template.py', encoding='utf-8') as f:
  20     MODULE_TEMPLATE = f.read()
  21
  22
  23 def main():
  24     parser = optparse.OptionParser(usage='%prog [OUTFILE.py]')
  25     args = parser.parse_args()[1] or ['yt_dlp/extractor/lazy_extractors.py']
  26     if len(args) != 1:
  27         parser.error('Expected only an output filename')
  28
  29     lazy_extractors_filename = args[0]
  30     if os.path.exists(lazy_extractors_filename):
  31         os.remove(lazy_extractors_filename)
  32
  33     _ALL_CLASSES = get_all_ies()  # Must be before import
  34
  35     from yt_dlp.extractor.common import InfoExtractor, SearchInfoExtractor
  36
  37     DummyInfoExtractor = type('InfoExtractor', (InfoExtractor,), {'IE_NAME': NO_ATTR})
  38     module_src = '\n'.join((
  39         MODULE_TEMPLATE,
  40         '    _module = None',
  41         *extra_ie_code(DummyInfoExtractor),
  42         '\nclass LazyLoadSearchExtractor(LazyLoadExtractor):\n    pass\n',
  43         *build_ies(_ALL_CLASSES, (InfoExtractor, SearchInfoExtractor), DummyInfoExtractor),
  44     ))
  45
  46     with open(lazy_extractors_filename, 'wt', encoding='utf-8') as f:
  47         f.write(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(PLUGINS_DIRNAME, BLOCKED_DIRNAME)
  55     try:
  56         from yt_dlp.extractor import _ALL_CLASSES
  57     finally:
  58         if os.path.exists(BLOCKED_DIRNAME):
  59             os.rename(BLOCKED_DIRNAME, PLUGINS_DIRNAME)
  60     return _ALL_CLASSES
  61
  62
  63 def extra_ie_code(ie, base=None):
  64     for var in STATIC_CLASS_PROPERTIES:
  65         val = getattr(ie, var)
  66         if val != (getattr(base, var) if base else NO_ATTR):
  67             yield f'    {var} = {val!r}'
  68     yield ''
  69
  70     for name in CLASS_METHODS:
  71         f = getattr(ie, name)
  72         if not base or f.__func__ != getattr(base, name).__func__:
  73             yield getsource(f)
  74
  75
  76 def build_ies(ies, bases, attr_base):
  77     names = []
  78     for ie in sort_ies(ies, bases):
  79         yield build_lazy_ie(ie, ie.__name__, attr_base)
  80         if ie in ies:
  81             names.append(ie.__name__)
  82
  83     yield f'\n_ALL_CLASSES = [{", ".join(names)}]'
  84
  85
  86 def sort_ies(ies, ignored_bases):
  87     """find the correct sorting and add the required base classes so that subclasses can be correctly created"""
  88     classes, returned_classes = ies[:-1], set()
  89     assert ies[-1].__name__ == 'GenericIE', 'Last IE must be GenericIE'
  90     while classes:
  91         for c in classes[:]:
  92             bases = set(c.__bases__) - {object, *ignored_bases}
  93             restart = False
  94             for b in bases:
  95                 if b not in classes and b not in returned_classes:
  96                     assert b.__name__ != 'GenericIE', 'Cannot inherit from GenericIE'
  97                     classes.insert(0, b)
  98                     restart = True
  99             if restart:
 100                 break
 101             if bases <= returned_classes:
 102                 yield c
 103                 returned_classes.add(c)
 104                 classes.remove(c)
 105                 break
 106     yield ies[-1]
 107
 108
 109 def build_lazy_ie(ie, name, attr_base):
 110     bases = ', '.join({
 111         'InfoExtractor': 'LazyLoadExtractor',
 112         'SearchInfoExtractor': 'LazyLoadSearchExtractor',
 113     }.get(base.__name__, base.__name__) for base in ie.__bases__)
 114
 115     s = IE_TEMPLATE.format(name=name, module=ie.__module__, bases=bases)
 116     valid_url = getattr(ie, '_VALID_URL', None)
 117     if not valid_url and hasattr(ie, '_make_valid_url'):
 118         valid_url = ie._make_valid_url()
 119     if valid_url:
 120         s += f'    _VALID_URL = {valid_url!r}\n'
 121     return s + '\n'.join(extra_ie_code(ie, attr_base))
 122
 123
 124 if __name__ == '__main__':
 125     main()