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