yt_dlp/postprocessor/common.py

   1 import functools
   2 import itertools
   3 import json
   4 import os
   5 import time
   6 import urllib.error
   7
   8 from ..utils import (
   9     PostProcessingError,
  10     _configuration_args,
  11     encodeFilename,
  12     network_exceptions,
  13     sanitized_Request,
  14     write_string,
  15 )
  16
  17
  18 class PostProcessorMetaClass(type):
  19     @staticmethod
  20     def run_wrapper(func):
  21         @functools.wraps(func)
  22         def run(self, info, *args, **kwargs):
  23             info_copy = self._copy_infodict(info)
  24             self._hook_progress({'status': 'started'}, info_copy)
  25             ret = func(self, info, *args, **kwargs)
  26             if ret is not None:
  27                 _, info = ret
  28             self._hook_progress({'status': 'finished'}, info_copy)
  29             return ret
  30         return run
  31
  32     def __new__(cls, name, bases, attrs):
  33         if 'run' in attrs:
  34             attrs['run'] = cls.run_wrapper(attrs['run'])
  35         return type.__new__(cls, name, bases, attrs)
  36
  37
  38 class PostProcessor(metaclass=PostProcessorMetaClass):
  39     """Post Processor class.
  40
  41     PostProcessor objects can be added to downloaders with their
  42     add_post_processor() method. When the downloader has finished a
  43     successful download, it will take its internal chain of PostProcessors
  44     and start calling the run() method on each one of them, first with
  45     an initial argument and then with the returned value of the previous
  46     PostProcessor.
  47
  48     The chain will be stopped if one of them ever returns None or the end
  49     of the chain is reached.
  50
  51     PostProcessor objects follow a "mutual registration" process similar
  52     to InfoExtractor objects.
  53
  54     Optionally PostProcessor can use a list of additional command-line arguments
  55     with self._configuration_args.
  56     """
  57
  58     _downloader = None
  59
  60     def __init__(self, downloader=None):
  61         self._progress_hooks = []
  62         self.add_progress_hook(self.report_progress)
  63         self.set_downloader(downloader)
  64         self.PP_NAME = self.pp_key()
  65
  66     @classmethod
  67     def pp_key(cls):
  68         name = cls.__name__[:-2]
  69         return name[6:] if name[:6].lower() == 'ffmpeg' else name
  70
  71     def to_screen(self, text, prefix=True, *args, **kwargs):
  72         tag = '[%s] ' % self.PP_NAME if prefix else ''
  73         if self._downloader:
  74             return self._downloader.to_screen(f'{tag}{text}', *args, **kwargs)
  75
  76     def report_warning(self, text, *args, **kwargs):
  77         if self._downloader:
  78             return self._downloader.report_warning(text, *args, **kwargs)
  79
  80     def deprecation_warning(self, text):
  81         if self._downloader:
  82             return self._downloader.deprecation_warning(text)
  83         write_string(f'DeprecationWarning: {text}')
  84
  85     def report_error(self, text, *args, **kwargs):
  86         self.deprecation_warning('"yt_dlp.postprocessor.PostProcessor.report_error" is deprecated. '
  87                                  'raise "yt_dlp.utils.PostProcessingError" instead')
  88         if self._downloader:
  89             return self._downloader.report_error(text, *args, **kwargs)
  90
  91     def write_debug(self, text, *args, **kwargs):
  92         if self._downloader:
  93             return self._downloader.write_debug(text, *args, **kwargs)
  94
  95     def get_param(self, name, default=None, *args, **kwargs):
  96         if self._downloader:
  97             return self._downloader.params.get(name, default, *args, **kwargs)
  98         return default
  99
 100     def set_downloader(self, downloader):
 101         """Sets the downloader for this PP."""
 102         self._downloader = downloader
 103         for ph in getattr(downloader, '_postprocessor_hooks', []):
 104             self.add_progress_hook(ph)
 105
 106     def _copy_infodict(self, info_dict):
 107         return getattr(self._downloader, '_copy_infodict', dict)(info_dict)
 108
 109     @staticmethod
 110     def _restrict_to(*, video=True, audio=True, images=True, simulated=True):
 111         allowed = {'video': video, 'audio': audio, 'images': images}
 112
 113         def decorator(func):
 114             @functools.wraps(func)
 115             def wrapper(self, info):
 116                 if not simulated and (self.get_param('simulate') or self.get_param('skip_download')):
 117                     return [], info
 118                 format_type = (
 119                     'video' if info.get('vcodec') != 'none'
 120                     else 'audio' if info.get('acodec') != 'none'
 121                     else 'images')
 122                 if allowed[format_type]:
 123                     return func(self, info)
 124                 else:
 125                     self.to_screen('Skipping %s' % format_type)
 126                     return [], info
 127             return wrapper
 128         return decorator
 129
 130     def run(self, information):
 131         """Run the PostProcessor.
 132
 133         The "information" argument is a dictionary like the ones
 134         composed by InfoExtractors. The only difference is that this
 135         one has an extra field called "filepath" that points to the
 136         downloaded file.
 137
 138         This method returns a tuple, the first element is a list of the files
 139         that can be deleted, and the second of which is the updated
 140         information.
 141
 142         In addition, this method may raise a PostProcessingError
 143         exception if post processing fails.
 144         """
 145         return [], information  # by default, keep file and do nothing
 146
 147     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 148         try:
 149             os.utime(encodeFilename(path), (atime, mtime))
 150         except Exception:
 151             self.report_warning(errnote)
 152
 153     def _configuration_args(self, exe, *args, **kwargs):
 154         return _configuration_args(
 155             self.pp_key(), self.get_param('postprocessor_args'), exe, *args, **kwargs)
 156
 157     def _hook_progress(self, status, info_dict):
 158         if not self._progress_hooks:
 159             return
 160         status.update({
 161             'info_dict': info_dict,
 162             'postprocessor': self.pp_key(),
 163         })
 164         for ph in self._progress_hooks:
 165             ph(status)
 166
 167     def add_progress_hook(self, ph):
 168         # See YoutubeDl.py (search for postprocessor_hooks) for a description of this interface
 169         self._progress_hooks.append(ph)
 170
 171     def report_progress(self, s):
 172         s['_default_template'] = '%(postprocessor)s %(status)s' % s
 173
 174         progress_dict = s.copy()
 175         progress_dict.pop('info_dict')
 176         progress_dict = {'info': s['info_dict'], 'progress': progress_dict}
 177
 178         progress_template = self.get_param('progress_template', {})
 179         tmpl = progress_template.get('postprocess')
 180         if tmpl:
 181             self._downloader.to_stdout(self._downloader.evaluate_outtmpl(tmpl, progress_dict))
 182
 183         self._downloader.to_console_title(self._downloader.evaluate_outtmpl(
 184             progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s',
 185             progress_dict))
 186
 187     def _download_json(self, url, *, expected_http_errors=(404,)):
 188         # While this is not an extractor, it behaves similar to one and
 189         # so obey extractor_retries and sleep_interval_requests
 190         max_retries = self.get_param('extractor_retries', 3)
 191         sleep_interval = self.get_param('sleep_interval_requests') or 0
 192
 193         self.write_debug(f'{self.PP_NAME} query: {url}')
 194         for retries in itertools.count():
 195             try:
 196                 rsp = self._downloader.urlopen(sanitized_Request(url))
 197                 return json.loads(rsp.read().decode(rsp.info().get_param('charset') or 'utf-8'))
 198             except network_exceptions as e:
 199                 if isinstance(e, urllib.error.HTTPError) and e.code in expected_http_errors:
 200                     return None
 201                 if retries < max_retries:
 202                     self.report_warning(f'{e}. Retrying...')
 203                     if sleep_interval > 0:
 204                         self.to_screen(f'Sleeping {sleep_interval} seconds ...')
 205                         time.sleep(sleep_interval)
 206                     continue
 207                 raise PostProcessingError(f'Unable to communicate with {self.PP_NAME} API: {e}')
 208
 209
 210 class AudioConversionError(PostProcessingError):
 211     pass