devscripts/make_lazy_extractors.py

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