yt_dlp/postprocessor/common.py

   1 from __future__ import unicode_literals
   2
   3 import functools
   4 import os
   5
   6 from ..compat import compat_str
   7 from ..utils import (
   8     _configuration_args,
   9     encodeFilename,
  10     PostProcessingError,
  11     write_string,
  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     The chain will be stopped if one of them ever returns None or the end
  46     of the chain is reached.
  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 compat_str(name[6:]) if name[:6].lower() == 'ffmpeg' else name
  67
  68     def to_screen(self, text, prefix=True, *args, **kwargs):
  69         tag = '[%s] ' % self.PP_NAME if prefix else ''
  70         if self._downloader:
  71             return self._downloader.to_screen('%s%s' % (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         # Exists only for compatibility. Do not use
  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 get_param(self, name, default=None, *args, **kwargs):
  92         if self._downloader:
  93             return self._downloader.params.get(name, default, *args, **kwargs)
  94         return default
  95
  96     def set_downloader(self, downloader):
  97         """Sets the downloader for this PP."""
  98         self._downloader = downloader
  99         for ph in getattr(downloader, '_postprocessor_hooks', []):
 100             self.add_progress_hook(ph)
 101
 102     def _copy_infodict(self, info_dict):
 103         return getattr(self._downloader, '_copy_infodict', dict)(info_dict)
 104
 105     @staticmethod
 106     def _restrict_to(*, video=True, audio=True, images=True, simulated=True):
 107         allowed = {'video': video, 'audio': audio, 'images': images}
 108
 109         def decorator(func):
 110             @functools.wraps(func)
 111             def wrapper(self, info):
 112                 if not simulated and (self.get_param('simulate') or self.get_param('skip_download')):
 113                     return [], info
 114                 format_type = (
 115                     'video' if info.get('vcodec') != 'none'
 116                     else 'audio' if info.get('acodec') != 'none'
 117                     else 'images')
 118                 if allowed[format_type]:
 119                     return func(self, info)
 120                 else:
 121                     self.to_screen('Skipping %s' % format_type)
 122                     return [], info
 123             return wrapper
 124         return decorator
 125
 126     def run(self, information):
 127         """Run the PostProcessor.
 128
 129         The "information" argument is a dictionary like the ones
 130         composed by InfoExtractors. The only difference is that this
 131         one has an extra field called "filepath" that points to the
 132         downloaded file.
 133
 134         This method returns a tuple, the first element is a list of the files
 135         that can be deleted, and the second of which is the updated
 136         information.
 137
 138         In addition, this method may raise a PostProcessingError
 139         exception if post processing fails.
 140         """
 141         return [], information  # by default, keep file and do nothing
 142
 143     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 144         try:
 145             os.utime(encodeFilename(path), (atime, mtime))
 146         except Exception:
 147             self.report_warning(errnote)
 148
 149     def _configuration_args(self, exe, *args, **kwargs):
 150         return _configuration_args(
 151             self.pp_key(), self.get_param('postprocessor_args'), exe, *args, **kwargs)
 152
 153     def _hook_progress(self, status, info_dict):
 154         if not self._progress_hooks:
 155             return
 156         status.update({
 157             'info_dict': info_dict,
 158             'postprocessor': self.pp_key(),
 159         })
 160         for ph in self._progress_hooks:
 161             ph(status)
 162
 163     def add_progress_hook(self, ph):
 164         # See YoutubeDl.py (search for postprocessor_hooks) for a description of this interface
 165         self._progress_hooks.append(ph)
 166
 167     def report_progress(self, s):
 168         s['_default_template'] = '%(postprocessor)s %(status)s' % s
 169
 170         progress_dict = s.copy()
 171         progress_dict.pop('info_dict')
 172         progress_dict = {'info': s['info_dict'], 'progress': progress_dict}
 173
 174         progress_template = self.get_param('progress_template', {})
 175         tmpl = progress_template.get('postprocess')
 176         if tmpl:
 177             self._downloader.to_stdout(self._downloader.evaluate_outtmpl(tmpl, progress_dict))
 178
 179         self._downloader.to_console_title(self._downloader.evaluate_outtmpl(
 180             progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s',
 181             progress_dict))
 182
 183
 184 class AudioConversionError(PostProcessingError):
 185     pass