yt_dlp/postprocessor/common.py

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