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