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     cli_configuration_args,
   9     encodeFilename,
  10     PostProcessingError,
  11 )
  12
  13
  14 class PostProcessor(object):
  15     """Post Processor class.
  16
  17     PostProcessor objects can be added to downloaders with their
  18     add_post_processor() method. When the downloader has finished a
  19     successful download, it will take its internal chain of PostProcessors
  20     and start calling the run() method on each one of them, first with
  21     an initial argument and then with the returned value of the previous
  22     PostProcessor.
  23
  24     The chain will be stopped if one of them ever returns None or the end
  25     of the chain is reached.
  26
  27     PostProcessor objects follow a "mutual registration" process similar
  28     to InfoExtractor objects.
  29
  30     Optionally PostProcessor can use a list of additional command-line arguments
  31     with self._configuration_args.
  32     """
  33
  34     _downloader = None
  35
  36     def __init__(self, downloader=None):
  37         self._downloader = downloader
  38         self.PP_NAME = self.pp_key()
  39
  40     @classmethod
  41     def pp_key(cls):
  42         name = cls.__name__[:-2]
  43         return compat_str(name[6:]) if name[:6].lower() == 'ffmpeg' else name
  44
  45     def to_screen(self, text, prefix=True, *args, **kwargs):
  46         tag = '[%s] ' % self.PP_NAME if prefix else ''
  47         if self._downloader:
  48             return self._downloader.to_screen('%s%s' % (tag, text), *args, **kwargs)
  49
  50     def report_warning(self, text, *args, **kwargs):
  51         if self._downloader:
  52             return self._downloader.report_warning(text, *args, **kwargs)
  53
  54     def report_error(self, text, *args, **kwargs):
  55         if self._downloader:
  56             return self._downloader.report_error(text, *args, **kwargs)
  57
  58     def write_debug(self, text, *args, **kwargs):
  59         if self._downloader:
  60             return self._downloader.write_debug(text, *args, **kwargs)
  61
  62     def get_param(self, name, default=None, *args, **kwargs):
  63         if self._downloader:
  64             return self._downloader.params.get(name, default, *args, **kwargs)
  65         return default
  66
  67     def set_downloader(self, downloader):
  68         """Sets the downloader for this PP."""
  69         self._downloader = downloader
  70
  71     @staticmethod
  72     def _restrict_to(*, video=True, audio=True, images=True):
  73         allowed = {'video': video, 'audio': audio, 'images': images}
  74
  75         def decorator(func):
  76             @functools.wraps(func)
  77             def wrapper(self, info):
  78                 format_type = (
  79                     'video' if info['vcodec'] != 'none'
  80                     else 'audio' if info['acodec'] != 'none'
  81                     else 'images')
  82                 if allowed[format_type]:
  83                     func(self, info)
  84                 else:
  85                     self.to_screen('Skipping %s' % format_type)
  86                     return [], info
  87             return wrapper
  88         return decorator
  89
  90     def run(self, information):
  91         """Run the PostProcessor.
  92
  93         The "information" argument is a dictionary like the ones
  94         composed by InfoExtractors. The only difference is that this
  95         one has an extra field called "filepath" that points to the
  96         downloaded file.
  97
  98         This method returns a tuple, the first element is a list of the files
  99         that can be deleted, and the second of which is the updated
 100         information.
 101
 102         In addition, this method may raise a PostProcessingError
 103         exception if post processing fails.
 104         """
 105         return [], information  # by default, keep file and do nothing
 106
 107     def try_utime(self, path, atime, mtime, errnote='Cannot update utime of file'):
 108         try:
 109             os.utime(encodeFilename(path), (atime, mtime))
 110         except Exception:
 111             self.report_warning(errnote)
 112
 113     def _configuration_args(self, exe, keys=None, default=[], use_compat=True):
 114         pp_key = self.pp_key().lower()
 115         exe = exe.lower()
 116         root_key = exe if pp_key == exe else '%s+%s' % (pp_key, exe)
 117         keys = ['%s%s' % (root_key, k) for k in (keys or [''])]
 118         if root_key in keys:
 119             keys += [root_key] + ([] if pp_key == exe else [(self.pp_key(), exe)]) + ['default']
 120         else:
 121             use_compat = False
 122         return cli_configuration_args(
 123             self.get_param('postprocessor_args'),
 124             keys, default, use_compat)
 125
 126
 127 class AudioConversionError(PostProcessingError):
 128     pass