yt_dlp/postprocessor/common.py

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