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