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     encodeFilename,
  11     network_exceptions,
  12     sanitized_Request,
  13     write_string,
  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, text):
  77         if self._downloader:
  78             return self._downloader.deprecation_warning(text)
  79         write_string(f'DeprecationWarning: {text}')
  80
  81     def report_error(self, text, *args, **kwargs):
  82         self.deprecation_warning('"yt_dlp.postprocessor.PostProcessor.report_error" is deprecated. '
  83                                  'raise "yt_dlp.utils.PostProcessingError" instead')
  84         if self._downloader:
  85             return self._downloader.report_error(text, *args, **kwargs)
  86
  87     def write_debug(self, text, *args, **kwargs):
  88         if self._downloader:
  89             return self._downloader.write_debug(text, *args, **kwargs)
  90
  91     def _delete_downloaded_files(self, *files_to_delete, **kwargs):
  92         if self._downloader:
  93             return self._downloader._delete_downloaded_files(*files_to_delete, **kwargs)
  94         for filename in set(filter(None, files_to_delete)):
  95             os.remove(filename)
  96
  97     def get_param(self, name, default=None, *args, **kwargs):
  98         if self._downloader:
  99             return self._downloader.params.get(name, default, *args, **kwargs)
 100         return default
 101
 102     def set_downloader(self, downloader):
 103         """Sets the downloader for this PP."""
 104         self._downloader = downloader
 105         for ph in getattr(downloader, '_postprocessor_hooks', []):
 106             self.add_progress_hook(ph)
 107
 108     def _copy_infodict(self, info_dict):
 109         return getattr(self._downloader, '_copy_infodict', dict)(info_dict)
 110
 111     @staticmethod
 112     def _restrict_to(*, video=True, audio=True, images=True, simulated=True):
 113         allowed = {'video': video, 'audio': audio, 'images': images}
 114
 115         def decorator(func):
 116             @functools.wraps(func)
 117             def wrapper(self, info):
 118                 if not simulated and (self.get_param('simulate') or self.get_param('skip_download')):
 119                     return [], info
 120                 format_type = (
 121                     'video' if info.get('vcodec') != 'none'
 122                     else 'audio' if info.get('acodec') != 'none'
 123                     else 'images')
 124                 if allowed[format_type]:
 125                     return func(self, info)
 126                 else:
 127                     self.to_screen('Skipping %s' % format_type)
 128                     return [], info
 129             return wrapper
 130         return decorator
 131
 132     def run(self, information):
 133         """Run the PostProcessor.
 134
 135         The "information" argument is a dictionary like the ones
 136         composed by InfoExtractors. The only difference is that this
 137         one has an extra field called "filepath" that points to the
 138         downloaded file.
 139
 140         This method returns a tuple, the first element is a list of the files
 141         that can be deleted, and the second of which is the updated
 142         information.
 143
 144         In addition, this method may raise a PostProcessingError
 145         exception if post processing fails.
 146         """
 147         return [], information  # by default, keep file and do nothing
 148
 149     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 150         try:
 151             os.utime(encodeFilename(path), (atime, mtime))
 152         except Exception:
 153             self.report_warning(errnote)
 154
 155     def _configuration_args(self, exe, *args, **kwargs):
 156         return _configuration_args(
 157             self.pp_key(), self.get_param('postprocessor_args'), exe, *args, **kwargs)
 158
 159     def _hook_progress(self, status, info_dict):
 160         if not self._progress_hooks:
 161             return
 162         status.update({
 163             'info_dict': info_dict,
 164             'postprocessor': self.pp_key(),
 165         })
 166         for ph in self._progress_hooks:
 167             ph(status)
 168
 169     def add_progress_hook(self, ph):
 170         # See YoutubeDl.py (search for postprocessor_hooks) for a description of this interface
 171         self._progress_hooks.append(ph)
 172
 173     def report_progress(self, s):
 174         s['_default_template'] = '%(postprocessor)s %(status)s' % s
 175         if not self._downloader:
 176             return
 177
 178         progress_dict = s.copy()
 179         progress_dict.pop('info_dict')
 180         progress_dict = {'info': s['info_dict'], 'progress': progress_dict}
 181
 182         progress_template = self.get_param('progress_template', {})
 183         tmpl = progress_template.get('postprocess')
 184         if tmpl:
 185             self._downloader.to_screen(
 186                 self._downloader.evaluate_outtmpl(tmpl, progress_dict), skip_eol=True, quiet=False)
 187
 188         self._downloader.to_console_title(self._downloader.evaluate_outtmpl(
 189             progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s',
 190             progress_dict))
 191
 192     def _retry_download(self, err, count, retries):
 193         # While this is not an extractor, it behaves similar to one and
 194         # so obey extractor_retries and sleep_interval_requests
 195         RetryManager.report_retry(err, count, retries, info=self.to_screen, warn=self.report_warning,
 196                                   sleep_func=self.get_param('sleep_interval_requests'))
 197
 198     def _download_json(self, url, *, expected_http_errors=(404,)):
 199         self.write_debug(f'{self.PP_NAME} query: {url}')
 200         for retry in RetryManager(self.get_param('extractor_retries', 3), self._retry_download):
 201             try:
 202                 rsp = self._downloader.urlopen(sanitized_Request(url))
 203             except network_exceptions as e:
 204                 if isinstance(e, urllib.error.HTTPError) and e.code in expected_http_errors:
 205                     return None
 206                 retry.error = PostProcessingError(f'Unable to communicate with {self.PP_NAME} API: {e}')
 207                 continue
 208         return json.loads(rsp.read().decode(rsp.info().get_param('charset') or 'utf-8'))
 209
 210
 211 class AudioConversionError(PostProcessingError):  # Deprecated
 212     pass