not relying on ffmpeg to do the post-processed file checking, instead doing it direct...
[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 AudioConversionError(BaseException):
61     def __init__(self, message):
62         self.message = message
63
64 class FFmpegExtractAudioPP(PostProcessor):
65     def __init__(self, downloader=None, preferredcodec=None, preferredquality=None, keepvideo=False, nopostoverwrites=False):
66         PostProcessor.__init__(self, downloader)
67         if preferredcodec is None:
68             preferredcodec = 'best'
69         self._preferredcodec = preferredcodec
70         self._preferredquality = preferredquality
71         self._keepvideo = keepvideo
72         self._nopostoverwrites = nopostoverwrites
73         self._exes = self.detect_executables()
74
75     @staticmethod
76     def detect_executables():
77         def executable(exe):
78             try:
79                 subprocess.Popen([exe, '-version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()
80             except OSError:
81                 return False
82             return exe
83         programs = ['avprobe', 'avconv', 'ffmpeg', 'ffprobe']
84         return dict((program, executable(program)) for program in programs)
85
86     def get_audio_codec(self, path):
87         if not self._exes['ffprobe'] and not self._exes['avprobe']: return None
88         try:
89             cmd = [self._exes['avprobe'] or self._exes['ffprobe'], '-show_streams', '--', encodeFilename(path)]
90             handle = subprocess.Popen(cmd, stderr=compat_subprocess_get_DEVNULL(), stdout=subprocess.PIPE)
91             output = handle.communicate()[0]
92             if handle.wait() != 0:
93                 return None
94         except (IOError, OSError):
95             return None
96         audio_codec = None
97         for line in output.decode('ascii', 'ignore').split('\n'):
98             if line.startswith('codec_name='):
99                 audio_codec = line.split('=')[1].strip()
100             elif line.strip() == 'codec_type=audio' and audio_codec is not None:
101                 return audio_codec
102         return None
103
104     def run_ffmpeg(self, path, out_path, codec, more_opts):
105         if not self._exes['ffmpeg'] and not self._exes['avconv']:
106             raise AudioConversionError('ffmpeg or avconv not found. Please install one.')
107         if codec is None:
108             acodec_opts = []
109         else:
110             acodec_opts = ['-acodec', codec]
111         cmd = ([self._exes['avconv'] or self._exes['ffmpeg'], '-y', '-i', encodeFilename(path), '-vn']
112                + acodec_opts + more_opts +
113                ['--', encodeFilename(out_path)])
114         p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
115         stdout,stderr = p.communicate()
116         if p.returncode != 0:
117             msg = stderr.strip().split('\n')[-1]
118             raise AudioConversionError(msg)
119
120     def run(self, information):
121         path = information['filepath']
122
123         filecodec = self.get_audio_codec(path)
124         if filecodec is None:
125             self._downloader.to_stderr(u'WARNING: unable to obtain file audio codec with ffprobe')
126             return None
127
128         more_opts = []
129         if self._preferredcodec == 'best' or self._preferredcodec == filecodec or (self._preferredcodec == 'm4a' and filecodec == 'aac'):
130             if self._preferredcodec == 'm4a' and filecodec == 'aac':
131                 # Lossless, but in another container
132                 acodec = 'copy'
133                 extension = self._preferredcodec
134                 more_opts = [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
135             elif filecodec in ['aac', 'mp3', 'vorbis']:
136                 # Lossless if possible
137                 acodec = 'copy'
138                 extension = filecodec
139                 if filecodec == 'aac':
140                     more_opts = ['-f', 'adts']
141                 if filecodec == 'vorbis':
142                     extension = 'ogg'
143             else:
144                 # MP3 otherwise.
145                 acodec = 'libmp3lame'
146                 extension = 'mp3'
147                 more_opts = []
148                 if self._preferredquality is not None:
149                     if int(self._preferredquality) < 10:
150                         more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
151                     else:
152                         more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
153         else:
154             # We convert the audio (lossy)
155             acodec = {'mp3': 'libmp3lame', 'aac': 'aac', 'm4a': 'aac', 'vorbis': 'libvorbis', 'wav': None}[self._preferredcodec]
156             extension = self._preferredcodec
157             more_opts = []
158             if self._preferredquality is not None:
159                 if int(self._preferredquality) < 10:
160                     more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
161                 else:
162                     more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
163             if self._preferredcodec == 'aac':
164                 more_opts += ['-f', 'adts']
165             if self._preferredcodec == 'm4a':
166                 more_opts += [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
167             if self._preferredcodec == 'vorbis':
168                 extension = 'ogg'
169             if self._preferredcodec == 'wav':
170                 extension = 'wav'
171                 more_opts += ['-f', 'wav']
172
173         prefix, sep, ext = path.rpartition(u'.') # not os.path.splitext, since the latter does not work on unicode in all setups
174         new_path = prefix + sep + extension
175         try:
176             if self._nopostoverwrites and os.path.exists(encodeFilename(new_path)):
177                 self._downloader.to_screen(u'[youtube] Post-process file %s exists, skipping' % new_path)
178             else:
179                 self._downloader.to_screen(u'[' + (self._exes['avconv'] and 'avconv' or 'ffmpeg') + '] Destination: ' + new_path)
180                 self.run_ffmpeg(path, new_path, acodec, more_opts)
181         except:
182             etype,e,tb = sys.exc_info()
183             if isinstance(e, AudioConversionError):
184                 self._downloader.to_stderr(u'ERROR: audio conversion failed: ' + e.message)
185             else:
186                 self._downloader.to_stderr(u'ERROR: error running ' + (self._exes['avconv'] and 'avconv' or 'ffmpeg'))
187             return None
188
189         # Try to update the date time for extracted audio file.
190         if information.get('filetime') is not None:
191             try:
192                 os.utime(encodeFilename(new_path), (time.time(), information['filetime']))
193             except:
194                 self._downloader.to_stderr(u'WARNING: Cannot update utime of audio file')
195
196         if not self._keepvideo:
197             try:
198                 os.remove(encodeFilename(path))
199             except (IOError, OSError):
200                 self._downloader.to_stderr(u'WARNING: Unable to remove downloaded video file')
201                 return None
202
203         information['filepath'] = new_path
204         return information