Add a PostProcessor for converting video format
[youtube-dl] / youtube_dl / PostProcessor.py
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3
4 from __future__ import absolute_import
5
6 import os
7 import subprocess
8 import sys
9 import time
10
11 from .utils import *
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
31     _downloader = None
32
33     def __init__(self, downloader=None):
34         self._downloader = downloader
35
36     def set_downloader(self, downloader):
37         """Sets the downloader for this PP."""
38         self._downloader = downloader
39
40     def run(self, information):
41         """Run the PostProcessor.
42
43         The "information" argument is a dictionary like the ones
44         composed by InfoExtractors. The only difference is that this
45         one has an extra field called "filepath" that points to the
46         downloaded file.
47
48         When this method returns None, the postprocessing chain is
49         stopped. However, this method may return an information
50         dictionary that will be passed to the next postprocessing
51         object in the chain. It can be the one it received after
52         changing some fields.
53
54         In addition, this method may raise a PostProcessingError
55         exception that will be taken into account by the downloader
56         it was called from.
57         """
58         return information # by default, do nothing
59
60 class FFmpegPostProcessorError(BaseException):
61     def __init__(self, message):
62         self.message = message
63
64 class AudioConversionError(BaseException):
65     def __init__(self, message):
66         self.message = message
67
68 class FFmpegPostProcessor(PostProcessor):
69     def __init__(self,downloader=None):
70         PostProcessor.__init__(self, downloader)
71         self._exes = self.detect_executables()
72
73     @staticmethod
74     def detect_executables():
75         def executable(exe):
76             try:
77                 subprocess.Popen([exe, '-version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()
78             except OSError:
79                 return False
80             return exe
81         programs = ['avprobe', 'avconv', 'ffmpeg', 'ffprobe']
82         return dict((program, executable(program)) for program in programs)
83
84     def run_ffmpeg(self, path, out_path, opts):
85         if not self._exes['ffmpeg'] and not self._exes['avconv']:
86             raise FFmpegPostProcessorError('ffmpeg or avconv not found. Please install one.')
87         cmd = ([self._exes['avconv'] or self._exes['ffmpeg'], '-y', '-i', encodeFilename(path)]
88                + opts +
89                [encodeFilename(self._ffmpeg_filename_argument(out_path))])
90         p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
91         stdout,stderr = p.communicate()
92         if p.returncode != 0:
93             msg = stderr.strip().split('\n')[-1]
94             raise FFmpegPostProcessorError(msg)
95
96     def _ffmpeg_filename_argument(self, fn):
97         # ffmpeg broke --, see https://ffmpeg.org/trac/ffmpeg/ticket/2127 for details
98         if fn.startswith(u'-'):
99             return u'./' + fn
100         return fn
101
102 class FFmpegExtractAudioPP(FFmpegPostProcessor):
103     def __init__(self, downloader=None, preferredcodec=None, preferredquality=None, keepvideo=False, nopostoverwrites=False):
104         FFmpegPostProcessor.__init__(self, downloader)
105         if preferredcodec is None:
106             preferredcodec = 'best'
107         self._preferredcodec = preferredcodec
108         self._preferredquality = preferredquality
109         self._keepvideo = keepvideo
110         self._nopostoverwrites = nopostoverwrites
111
112     def get_audio_codec(self, path):
113         if not self._exes['ffprobe'] and not self._exes['avprobe']: return None
114         try:
115             cmd = [self._exes['avprobe'] or self._exes['ffprobe'], '-show_streams', encodeFilename(self._ffmpeg_filename_argument(path))]
116             handle = subprocess.Popen(cmd, stderr=compat_subprocess_get_DEVNULL(), stdout=subprocess.PIPE)
117             output = handle.communicate()[0]
118             if handle.wait() != 0:
119                 return None
120         except (IOError, OSError):
121             return None
122         audio_codec = None
123         for line in output.decode('ascii', 'ignore').split('\n'):
124             if line.startswith('codec_name='):
125                 audio_codec = line.split('=')[1].strip()
126             elif line.strip() == 'codec_type=audio' and audio_codec is not None:
127                 return audio_codec
128         return None
129
130     def run_ffmpeg(self, path, out_path, codec, more_opts):
131         if not self._exes['ffmpeg'] and not self._exes['avconv']:
132             raise AudioConversionError('ffmpeg or avconv not found. Please install one.')
133         if codec is None:
134             acodec_opts = []
135         else:
136             acodec_opts = ['-acodec', codec]
137         opts = ['-vn'] + acodec_opts + more_opts
138         try:
139             FFmpegPostProcessor.run_ffmpeg(self, path, out_path, opts)
140         except FFmpegPostProcessorError as err:
141             raise AudioConversionError(err.message)
142
143     def run(self, information):
144         path = information['filepath']
145
146         filecodec = self.get_audio_codec(path)
147         if filecodec is None:
148             self._downloader.to_stderr(u'WARNING: unable to obtain file audio codec with ffprobe')
149             return None
150
151         more_opts = []
152         if self._preferredcodec == 'best' or self._preferredcodec == filecodec or (self._preferredcodec == 'm4a' and filecodec == 'aac'):
153             if self._preferredcodec == 'm4a' and filecodec == 'aac':
154                 # Lossless, but in another container
155                 acodec = 'copy'
156                 extension = self._preferredcodec
157                 more_opts = [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
158             elif filecodec in ['aac', 'mp3', 'vorbis', 'opus']:
159                 # Lossless if possible
160                 acodec = 'copy'
161                 extension = filecodec
162                 if filecodec == 'aac':
163                     more_opts = ['-f', 'adts']
164                 if filecodec == 'vorbis':
165                     extension = 'ogg'
166             else:
167                 # MP3 otherwise.
168                 acodec = 'libmp3lame'
169                 extension = 'mp3'
170                 more_opts = []
171                 if self._preferredquality is not None:
172                     if int(self._preferredquality) < 10:
173                         more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
174                     else:
175                         more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
176         else:
177             # We convert the audio (lossy)
178             acodec = {'mp3': 'libmp3lame', 'aac': 'aac', 'm4a': 'aac', 'opus': 'opus', 'vorbis': 'libvorbis', 'wav': None}[self._preferredcodec]
179             extension = self._preferredcodec
180             more_opts = []
181             if self._preferredquality is not None:
182                 if int(self._preferredquality) < 10:
183                     more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
184                 else:
185                     more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
186             if self._preferredcodec == 'aac':
187                 more_opts += ['-f', 'adts']
188             if self._preferredcodec == 'm4a':
189                 more_opts += [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
190             if self._preferredcodec == 'vorbis':
191                 extension = 'ogg'
192             if self._preferredcodec == 'wav':
193                 extension = 'wav'
194                 more_opts += ['-f', 'wav']
195
196         prefix, sep, ext = path.rpartition(u'.') # not os.path.splitext, since the latter does not work on unicode in all setups
197         new_path = prefix + sep + extension
198         try:
199             if self._nopostoverwrites and os.path.exists(encodeFilename(new_path)):
200                 self._downloader.to_screen(u'[youtube] Post-process file %s exists, skipping' % new_path)
201             else:
202                 self._downloader.to_screen(u'[' + (self._exes['avconv'] and 'avconv' or 'ffmpeg') + '] Destination: ' + new_path)
203                 self.run_ffmpeg(path, new_path, acodec, more_opts)
204         except:
205             etype,e,tb = sys.exc_info()
206             if isinstance(e, AudioConversionError):
207                 self._downloader.to_stderr(u'ERROR: audio conversion failed: ' + e.message)
208             else:
209                 self._downloader.to_stderr(u'ERROR: error running ' + (self._exes['avconv'] and 'avconv' or 'ffmpeg'))
210             return None
211
212         # Try to update the date time for extracted audio file.
213         if information.get('filetime') is not None:
214             try:
215                 os.utime(encodeFilename(new_path), (time.time(), information['filetime']))
216             except:
217                 self._downloader.to_stderr(u'WARNING: Cannot update utime of audio file')
218
219         if not self._keepvideo:
220             try:
221                 os.remove(encodeFilename(path))
222             except (IOError, OSError):
223                 self._downloader.to_stderr(u'WARNING: Unable to remove downloaded video file')
224                 return None
225
226         information['filepath'] = new_path
227         return information
228
229 class FFmpegVideoConvertor(FFmpegPostProcessor):
230     def __init__(self, downloader=None,preferedformat=None):
231         FFmpegPostProcessor.__init__(self,downloader)
232         self._preferedformat=preferedformat
233
234     def run(self, information):
235         path = information['filepath']
236         prefix, sep, ext = path.rpartition(u'.')
237         outpath = prefix + sep + self._preferedformat
238         if not self._preferedformat or information['format'] == self._preferedformat:
239             return information
240         self._downloader.to_screen(u'['+'ffmpeg'+'] Converting video from %s to %s, Destination: ' % (information['format'], self._preferedformat) +outpath)
241         self.run_ffmpeg(path, outpath, [])
242         information['filepath'] = outpath
243         information['format'] = self._preferedformat
244         return information