yt_dlp/postprocessor/common.py

   1 from __future__ import unicode_literals
   2
   3 import copy
   4 import functools
   5 import os
   6
   7 from ..compat import compat_str
   8 from ..utils import (
   9     _configuration_args,
  10     encodeFilename,
  11     PostProcessingError,
  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             self._hook_progress({'status': 'started'}, info)
  21             ret = func(self, info, *args, **kwargs)
  22             if ret is not None:
  23                 _, info = ret
  24             self._hook_progress({'status': 'finished'}, info)
  25             return ret
  26         return run
  27
  28     def __new__(cls, name, bases, attrs):
  29         if 'run' in attrs:
  30             attrs['run'] = cls.run_wrapper(attrs['run'])
  31         return type.__new__(cls, name, bases, attrs)
  32
  33
  34 class PostProcessor(metaclass=PostProcessorMetaClass):
  35     """Post Processor class.
  36
  37     PostProcessor objects can be added to downloaders with their
  38     add_post_processor() method. When the downloader has finished a
  39     successful download, it will take its internal chain of PostProcessors
  40     and start calling the run() method on each one of them, first with
  41     an initial argument and then with the returned value of the previous
  42     PostProcessor.
  43
  44     The chain will be stopped if one of them ever returns None or the end
  45     of the chain is reached.
  46
  47     PostProcessor objects follow a "mutual registration" process similar
  48     to InfoExtractor objects.
  49
  50     Optionally PostProcessor can use a list of additional command-line arguments
  51     with self._configuration_args.
  52     """
  53
  54     _downloader = None
  55
  56     def __init__(self, downloader=None):
  57         self._progress_hooks = []
  58         self.add_progress_hook(self.report_progress)
  59         self.set_downloader(downloader)
  60         self.PP_NAME = self.pp_key()
  61
  62     @classmethod
  63     def pp_key(cls):
  64         name = cls.__name__[:-2]
  65         return compat_str(name[6:]) if name[:6].lower() == 'ffmpeg' else name
  66
  67     def to_screen(self, text, prefix=True, *args, **kwargs):
  68         tag = '[%s] ' % self.PP_NAME if prefix else ''
  69         if self._downloader:
  70             return self._downloader.to_screen('%s%s' % (tag, text), *args, **kwargs)
  71
  72     def report_warning(self, text, *args, **kwargs):
  73         if self._downloader:
  74             return self._downloader.report_warning(text, *args, **kwargs)
  75
  76     def report_error(self, text, *args, **kwargs):
  77         # Exists only for compatibility. Do not use
  78         if self._downloader:
  79             return self._downloader.report_error(text, *args, **kwargs)
  80
  81     def write_debug(self, text, *args, **kwargs):
  82         if self._downloader:
  83             return self._downloader.write_debug(text, *args, **kwargs)
  84
  85     def get_param(self, name, default=None, *args, **kwargs):
  86         if self._downloader:
  87             return self._downloader.params.get(name, default, *args, **kwargs)
  88         return default
  89
  90     def set_downloader(self, downloader):
  91         """Sets the downloader for this PP."""
  92         self._downloader = downloader
  93         if not downloader:
  94             return
  95         for ph in downloader._postprocessor_hooks:
  96             self.add_progress_hook(ph)
  97
  98     @staticmethod
  99     def _restrict_to(*, video=True, audio=True, images=True):
 100         allowed = {'video': video, 'audio': audio, 'images': images}
 101
 102         def decorator(func):
 103             @functools.wraps(func)
 104             def wrapper(self, info):
 105                 format_type = (
 106                     'video' if info.get('vcodec') != 'none'
 107                     else 'audio' if info.get('acodec') != 'none'
 108                     else 'images')
 109                 if allowed[format_type]:
 110                     return func(self, info)
 111                 else:
 112                     self.to_screen('Skipping %s' % format_type)
 113                     return [], info
 114             return wrapper
 115         return decorator
 116
 117     def run(self, information):
 118         """Run the PostProcessor.
 119
 120         The "information" argument is a dictionary like the ones
 121         composed by InfoExtractors. The only difference is that this
 122         one has an extra field called "filepath" that points to the
 123         downloaded file.
 124
 125         This method returns a tuple, the first element is a list of the files
 126         that can be deleted, and the second of which is the updated
 127         information.
 128
 129         In addition, this method may raise a PostProcessingError
 130         exception if post processing fails.
 131         """
 132         return [], information  # by default, keep file and do nothing
 133
 134     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 135         try:
 136             os.utime(encodeFilename(path), (atime, mtime))
 137         except Exception:
 138             self.report_warning(errnote)
 139
 140     def _configuration_args(self, exe, *args, **kwargs):
 141         return _configuration_args(
 142             self.pp_key(), self.get_param('postprocessor_args'), exe, *args, **kwargs)
 143
 144     def _hook_progress(self, status, info_dict):
 145         if not self._progress_hooks:
 146             return
 147         info_dict = dict(info_dict)
 148         for key in ('__original_infodict', '__postprocessors'):
 149             info_dict.pop(key, None)
 150         status.update({
 151             'info_dict': copy.deepcopy(info_dict),
 152             'postprocessor': self.pp_key(),
 153         })
 154         for ph in self._progress_hooks:
 155             ph(status)
 156
 157     def add_progress_hook(self, ph):
 158         # See YoutubeDl.py (search for postprocessor_hooks) for a description of this interface
 159         self._progress_hooks.append(ph)
 160
 161     def report_progress(self, s):
 162         s['_default_template'] = '%(postprocessor)s %(status)s' % s
 163
 164         progress_dict = s.copy()
 165         progress_dict.pop('info_dict')
 166         progress_dict = {'info': s['info_dict'], 'progress': progress_dict}
 167
 168         progress_template = self.get_param('progress_template', {})
 169         tmpl = progress_template.get('postprocess')
 170         if tmpl:
 171             self._downloader.to_stdout(self._downloader.evaluate_outtmpl(tmpl, progress_dict))
 172
 173         self._downloader.to_console_title(self._downloader.evaluate_outtmpl(
 174             progress_template.get('postprocess-title') or 'yt-dlp %(progress._default_template)s',
 175             progress_dict))
 176
 177
 178 class AudioConversionError(PostProcessingError):
 179     pass