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):
 107         allowed = {'video': video, 'audio': audio, 'images': images}
 108
 109         def decorator(func):
 110             @functools.wraps(func)
 111             def wrapper(self, info):
 112                 format_type = (
 113                     'video' if info.get('vcodec') != 'none'
 114                     else 'audio' if info.get('acodec') != 'none'
 115                     else 'images')
 116                 if allowed[format_type]:
 117                     return func(self, info)
 118                 else:
 119                     self.to_screen('Skipping %s' % format_type)
 120                     return [], info
 121             return wrapper
 122         return decorator
 123
 124     def run(self, information):
 125         """Run the PostProcessor.
 126
 127         The "information" argument is a dictionary like the ones
 128         composed by InfoExtractors. The only difference is that this
 129         one has an extra field called "filepath" that points to the
 130         downloaded file.
 131
 132         This method returns a tuple, the first element is a list of the files
 133         that can be deleted, and the second of which is the updated
 134         information.
 135
 136         In addition, this method may raise a PostProcessingError
 137         exception if post processing fails.
 138         """
 139         return [], information  # by default, keep file and do nothing
 140
 141     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 142         try:
 143             os.utime(encodeFilename(path), (atime, mtime))
 144         except Exception:
 145             self.report_warning(errnote)
 146
 147     def _configuration_args(self, exe, *args, **kwargs):
 148         return _configuration_args(
 149             self.pp_key(), self.get_param('postprocessor_args'), exe, *args, **kwargs)
 150
 151     def _hook_progress(self, status, info_dict):
 152         if not self._progress_hooks:
 153             return
 154         status.update({
 155             'info_dict': info_dict,
 156             'postprocessor': self.pp_key(),
 157         })
 158         for ph in self._progress_hooks:
 159             ph(status)
 160
 161     def add_progress_hook(self, ph):
 162         # See YoutubeDl.py (search for postprocessor_hooks) for a description of this interface
 163         self._progress_hooks.append(ph)
 164
 165     def report_progress(self, s):
 166         s['_default_template'] = '%(postprocessor)s %(status)s' % s
 167
 168         progress_dict = s.copy()
 169         progress_dict.pop('info_dict')
 170         progress_dict = {'info': s['info_dict'], 'progress': progress_dict}
 171
 172         progress_template = self.get_param('progress_template', {})
 173         tmpl = progress_template.get('postprocess')
 174         if tmpl:
 175             self._downloader.to_stdout(self._downloader.evaluate_outtmpl(tmpl, progress_dict))
 176
 177         self._downloader.to_console_title(self._downloader.evaluate_outtmpl(
 178             progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s',
 179             progress_dict))
 180
 181
 182 class AudioConversionError(PostProcessingError):
 183     pass