Move postprocessor into its own package
[youtube-dl] / youtube_dl / postprocessor / __init__.py
1 import os
2 import subprocess
3 import sys
4 import time
5
6
7 from ..utils import (
8     compat_subprocess_get_DEVNULL,
9     encodeFilename,
10     hyphenate_date,
11     PostProcessingError,
12     prepend_extension,
13     shell_quote,
14     subtitles_filename,
15 )
16
17
18 class PostProcessor(object):
19     """Post Processor class.
20
21     PostProcessor objects can be added to downloaders with their
22     add_post_processor() method. When the downloader has finished a
23     successful download, it will take its internal chain of PostProcessors
24     and start calling the run() method on each one of them, first with
25     an initial argument and then with the returned value of the previous
26     PostProcessor.
27
28     The chain will be stopped if one of them ever returns None or the end
29     of the chain is reached.
30
31     PostProcessor objects follow a "mutual registration" process similar
32     to InfoExtractor objects.
33     """
34
35     _downloader = None
36
37     def __init__(self, downloader=None):
38         self._downloader = downloader
39
40     def set_downloader(self, downloader):
41         """Sets the downloader for this PP."""
42         self._downloader = downloader
43
44     def run(self, information):
45         """Run the PostProcessor.
46
47         The "information" argument is a dictionary like the ones
48         composed by InfoExtractors. The only difference is that this
49         one has an extra field called "filepath" that points to the
50         downloaded file.
51
52         This method returns a tuple, the first element of which describes
53         whether the original file should be kept (i.e. not deleted - None for
54         no preference), and the second of which is the updated information.
55
56         In addition, this method may raise a PostProcessingError
57         exception if post processing fails.
58         """
59         return None, information # by default, keep file and do nothing
60
61 class FFmpegPostProcessorError(PostProcessingError):
62     pass
63
64 class AudioConversionError(PostProcessingError):
65     pass
66
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_multiple_files(self, input_paths, out_path, opts):
85         if not self._exes['ffmpeg'] and not self._exes['avconv']:
86             raise FFmpegPostProcessorError(u'ffmpeg or avconv not found. Please install one.')
87
88         files_cmd = []
89         for path in input_paths:
90             files_cmd.extend(['-i', encodeFilename(path, True)])
91         cmd = ([self._exes['avconv'] or self._exes['ffmpeg'], '-y'] + files_cmd
92                + opts +
93                [encodeFilename(self._ffmpeg_filename_argument(out_path), True)])
94
95         if self._downloader.params.get('verbose', False):
96             self._downloader.to_screen(u'[debug] ffmpeg command line: %s' % shell_quote(cmd))
97         p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
98         stdout,stderr = p.communicate()
99         if p.returncode != 0:
100             stderr = stderr.decode('utf-8', 'replace')
101             msg = stderr.strip().split('\n')[-1]
102             raise FFmpegPostProcessorError(msg)
103
104     def run_ffmpeg(self, path, out_path, opts):
105         self.run_ffmpeg_multiple_files([path], out_path, opts)
106
107     def _ffmpeg_filename_argument(self, fn):
108         # ffmpeg broke --, see https://ffmpeg.org/trac/ffmpeg/ticket/2127 for details
109         if fn.startswith(u'-'):
110             return u'./' + fn
111         return fn
112
113
114 class FFmpegExtractAudioPP(FFmpegPostProcessor):
115     def __init__(self, downloader=None, preferredcodec=None, preferredquality=None, nopostoverwrites=False):
116         FFmpegPostProcessor.__init__(self, downloader)
117         if preferredcodec is None:
118             preferredcodec = 'best'
119         self._preferredcodec = preferredcodec
120         self._preferredquality = preferredquality
121         self._nopostoverwrites = nopostoverwrites
122
123     def get_audio_codec(self, path):
124         if not self._exes['ffprobe'] and not self._exes['avprobe']:
125             raise PostProcessingError(u'ffprobe or avprobe not found. Please install one.')
126         try:
127             cmd = [
128                 self._exes['avprobe'] or self._exes['ffprobe'],
129                 '-show_streams',
130                 encodeFilename(self._ffmpeg_filename_argument(path), True)]
131             handle = subprocess.Popen(cmd, stderr=compat_subprocess_get_DEVNULL(), stdout=subprocess.PIPE)
132             output = handle.communicate()[0]
133             if handle.wait() != 0:
134                 return None
135         except (IOError, OSError):
136             return None
137         audio_codec = None
138         for line in output.decode('ascii', 'ignore').split('\n'):
139             if line.startswith('codec_name='):
140                 audio_codec = line.split('=')[1].strip()
141             elif line.strip() == 'codec_type=audio' and audio_codec is not None:
142                 return audio_codec
143         return None
144
145     def run_ffmpeg(self, path, out_path, codec, more_opts):
146         if not self._exes['ffmpeg'] and not self._exes['avconv']:
147             raise AudioConversionError('ffmpeg or avconv not found. Please install one.')
148         if codec is None:
149             acodec_opts = []
150         else:
151             acodec_opts = ['-acodec', codec]
152         opts = ['-vn'] + acodec_opts + more_opts
153         try:
154             FFmpegPostProcessor.run_ffmpeg(self, path, out_path, opts)
155         except FFmpegPostProcessorError as err:
156             raise AudioConversionError(err.msg)
157
158     def run(self, information):
159         path = information['filepath']
160
161         filecodec = self.get_audio_codec(path)
162         if filecodec is None:
163             raise PostProcessingError(u'WARNING: unable to obtain file audio codec with ffprobe')
164
165         more_opts = []
166         if self._preferredcodec == 'best' or self._preferredcodec == filecodec or (self._preferredcodec == 'm4a' and filecodec == 'aac'):
167             if filecodec == 'aac' and self._preferredcodec in ['m4a', 'best']:
168                 # Lossless, but in another container
169                 acodec = 'copy'
170                 extension = 'm4a'
171                 more_opts = [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
172             elif filecodec in ['aac', 'mp3', 'vorbis', 'opus']:
173                 # Lossless if possible
174                 acodec = 'copy'
175                 extension = filecodec
176                 if filecodec == 'aac':
177                     more_opts = ['-f', 'adts']
178                 if filecodec == 'vorbis':
179                     extension = 'ogg'
180             else:
181                 # MP3 otherwise.
182                 acodec = 'libmp3lame'
183                 extension = 'mp3'
184                 more_opts = []
185                 if self._preferredquality is not None:
186                     if int(self._preferredquality) < 10:
187                         more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
188                     else:
189                         more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
190         else:
191             # We convert the audio (lossy)
192             acodec = {'mp3': 'libmp3lame', 'aac': 'aac', 'm4a': 'aac', 'opus': 'opus', 'vorbis': 'libvorbis', 'wav': None}[self._preferredcodec]
193             extension = self._preferredcodec
194             more_opts = []
195             if self._preferredquality is not None:
196                 # The opus codec doesn't support the -aq option
197                 if int(self._preferredquality) < 10 and extension != 'opus':
198                     more_opts += [self._exes['avconv'] and '-q:a' or '-aq', self._preferredquality]
199                 else:
200                     more_opts += [self._exes['avconv'] and '-b:a' or '-ab', self._preferredquality + 'k']
201             if self._preferredcodec == 'aac':
202                 more_opts += ['-f', 'adts']
203             if self._preferredcodec == 'm4a':
204                 more_opts += [self._exes['avconv'] and '-bsf:a' or '-absf', 'aac_adtstoasc']
205             if self._preferredcodec == 'vorbis':
206                 extension = 'ogg'
207             if self._preferredcodec == 'wav':
208                 extension = 'wav'
209                 more_opts += ['-f', 'wav']
210
211         prefix, sep, ext = path.rpartition(u'.') # not os.path.splitext, since the latter does not work on unicode in all setups
212         new_path = prefix + sep + extension
213
214         # If we download foo.mp3 and convert it to... foo.mp3, then don't delete foo.mp3, silly.
215         if new_path == path:
216             self._nopostoverwrites = True
217
218         try:
219             if self._nopostoverwrites and os.path.exists(encodeFilename(new_path)):
220                 self._downloader.to_screen(u'[youtube] Post-process file %s exists, skipping' % new_path)
221             else:
222                 self._downloader.to_screen(u'[' + (self._exes['avconv'] and 'avconv' or 'ffmpeg') + '] Destination: ' + new_path)
223                 self.run_ffmpeg(path, new_path, acodec, more_opts)
224         except:
225             etype,e,tb = sys.exc_info()
226             if isinstance(e, AudioConversionError):
227                 msg = u'audio conversion failed: ' + e.msg
228             else:
229                 msg = u'error running ' + (self._exes['avconv'] and 'avconv' or 'ffmpeg')
230             raise PostProcessingError(msg)
231
232         # Try to update the date time for extracted audio file.
233         if information.get('filetime') is not None:
234             try:
235                 os.utime(encodeFilename(new_path), (time.time(), information['filetime']))
236             except:
237                 self._downloader.report_warning(u'Cannot update utime of audio file')
238
239         information['filepath'] = new_path
240         return self._nopostoverwrites,information
241
242
243 class FFmpegVideoConvertor(FFmpegPostProcessor):
244     def __init__(self, downloader=None,preferedformat=None):
245         super(FFmpegVideoConvertor, self).__init__(downloader)
246         self._preferedformat=preferedformat
247
248     def run(self, information):
249         path = information['filepath']
250         prefix, sep, ext = path.rpartition(u'.')
251         outpath = prefix + sep + self._preferedformat
252         if information['ext'] == self._preferedformat:
253             self._downloader.to_screen(u'[ffmpeg] Not converting video file %s - already is in target format %s' % (path, self._preferedformat))
254             return True,information
255         self._downloader.to_screen(u'['+'ffmpeg'+'] Converting video from %s to %s, Destination: ' % (information['ext'], self._preferedformat) +outpath)
256         self.run_ffmpeg(path, outpath, [])
257         information['filepath'] = outpath
258         information['format'] = self._preferedformat
259         information['ext'] = self._preferedformat
260         return False,information
261
262
263 class FFmpegEmbedSubtitlePP(FFmpegPostProcessor):
264     # See http://www.loc.gov/standards/iso639-2/ISO-639-2_utf-8.txt
265     _lang_map = {
266         'aa': 'aar',
267         'ab': 'abk',
268         'ae': 'ave',
269         'af': 'afr',
270         'ak': 'aka',
271         'am': 'amh',
272         'an': 'arg',
273         'ar': 'ara',
274         'as': 'asm',
275         'av': 'ava',
276         'ay': 'aym',
277         'az': 'aze',
278         'ba': 'bak',
279         'be': 'bel',
280         'bg': 'bul',
281         'bh': 'bih',
282         'bi': 'bis',
283         'bm': 'bam',
284         'bn': 'ben',
285         'bo': 'bod',
286         'br': 'bre',
287         'bs': 'bos',
288         'ca': 'cat',
289         'ce': 'che',
290         'ch': 'cha',
291         'co': 'cos',
292         'cr': 'cre',
293         'cs': 'ces',
294         'cu': 'chu',
295         'cv': 'chv',
296         'cy': 'cym',
297         'da': 'dan',
298         'de': 'deu',
299         'dv': 'div',
300         'dz': 'dzo',
301         'ee': 'ewe',
302         'el': 'ell',
303         'en': 'eng',
304         'eo': 'epo',
305         'es': 'spa',
306         'et': 'est',
307         'eu': 'eus',
308         'fa': 'fas',
309         'ff': 'ful',
310         'fi': 'fin',
311         'fj': 'fij',
312         'fo': 'fao',
313         'fr': 'fra',
314         'fy': 'fry',
315         'ga': 'gle',
316         'gd': 'gla',
317         'gl': 'glg',
318         'gn': 'grn',
319         'gu': 'guj',
320         'gv': 'glv',
321         'ha': 'hau',
322         'he': 'heb',
323         'hi': 'hin',
324         'ho': 'hmo',
325         'hr': 'hrv',
326         'ht': 'hat',
327         'hu': 'hun',
328         'hy': 'hye',
329         'hz': 'her',
330         'ia': 'ina',
331         'id': 'ind',
332         'ie': 'ile',
333         'ig': 'ibo',
334         'ii': 'iii',
335         'ik': 'ipk',
336         'io': 'ido',
337         'is': 'isl',
338         'it': 'ita',
339         'iu': 'iku',
340         'ja': 'jpn',
341         'jv': 'jav',
342         'ka': 'kat',
343         'kg': 'kon',
344         'ki': 'kik',
345         'kj': 'kua',
346         'kk': 'kaz',
347         'kl': 'kal',
348         'km': 'khm',
349         'kn': 'kan',
350         'ko': 'kor',
351         'kr': 'kau',
352         'ks': 'kas',
353         'ku': 'kur',
354         'kv': 'kom',
355         'kw': 'cor',
356         'ky': 'kir',
357         'la': 'lat',
358         'lb': 'ltz',
359         'lg': 'lug',
360         'li': 'lim',
361         'ln': 'lin',
362         'lo': 'lao',
363         'lt': 'lit',
364         'lu': 'lub',
365         'lv': 'lav',
366         'mg': 'mlg',
367         'mh': 'mah',
368         'mi': 'mri',
369         'mk': 'mkd',
370         'ml': 'mal',
371         'mn': 'mon',
372         'mr': 'mar',
373         'ms': 'msa',
374         'mt': 'mlt',
375         'my': 'mya',
376         'na': 'nau',
377         'nb': 'nob',
378         'nd': 'nde',
379         'ne': 'nep',
380         'ng': 'ndo',
381         'nl': 'nld',
382         'nn': 'nno',
383         'no': 'nor',
384         'nr': 'nbl',
385         'nv': 'nav',
386         'ny': 'nya',
387         'oc': 'oci',
388         'oj': 'oji',
389         'om': 'orm',
390         'or': 'ori',
391         'os': 'oss',
392         'pa': 'pan',
393         'pi': 'pli',
394         'pl': 'pol',
395         'ps': 'pus',
396         'pt': 'por',
397         'qu': 'que',
398         'rm': 'roh',
399         'rn': 'run',
400         'ro': 'ron',
401         'ru': 'rus',
402         'rw': 'kin',
403         'sa': 'san',
404         'sc': 'srd',
405         'sd': 'snd',
406         'se': 'sme',
407         'sg': 'sag',
408         'si': 'sin',
409         'sk': 'slk',
410         'sl': 'slv',
411         'sm': 'smo',
412         'sn': 'sna',
413         'so': 'som',
414         'sq': 'sqi',
415         'sr': 'srp',
416         'ss': 'ssw',
417         'st': 'sot',
418         'su': 'sun',
419         'sv': 'swe',
420         'sw': 'swa',
421         'ta': 'tam',
422         'te': 'tel',
423         'tg': 'tgk',
424         'th': 'tha',
425         'ti': 'tir',
426         'tk': 'tuk',
427         'tl': 'tgl',
428         'tn': 'tsn',
429         'to': 'ton',
430         'tr': 'tur',
431         'ts': 'tso',
432         'tt': 'tat',
433         'tw': 'twi',
434         'ty': 'tah',
435         'ug': 'uig',
436         'uk': 'ukr',
437         'ur': 'urd',
438         'uz': 'uzb',
439         've': 'ven',
440         'vi': 'vie',
441         'vo': 'vol',
442         'wa': 'wln',
443         'wo': 'wol',
444         'xh': 'xho',
445         'yi': 'yid',
446         'yo': 'yor',
447         'za': 'zha',
448         'zh': 'zho',
449         'zu': 'zul',
450     }
451
452     def __init__(self, downloader=None, subtitlesformat='srt'):
453         super(FFmpegEmbedSubtitlePP, self).__init__(downloader)
454         self._subformat = subtitlesformat
455
456     @classmethod
457     def _conver_lang_code(cls, code):
458         """Convert language code from ISO 639-1 to ISO 639-2/T"""
459         return cls._lang_map.get(code[:2])
460
461     def run(self, information):
462         if information['ext'] != u'mp4':
463             self._downloader.to_screen(u'[ffmpeg] Subtitles can only be embedded in mp4 files')
464             return True, information
465         if not information.get('subtitles'):
466             self._downloader.to_screen(u'[ffmpeg] There aren\'t any subtitles to embed') 
467             return True, information
468
469         sub_langs = [key for key in information['subtitles']]
470         filename = information['filepath']
471         input_files = [filename] + [subtitles_filename(filename, lang, self._subformat) for lang in sub_langs]
472
473         opts = ['-map', '0:0', '-map', '0:1', '-c:v', 'copy', '-c:a', 'copy']
474         for (i, lang) in enumerate(sub_langs):
475             opts.extend(['-map', '%d:0' % (i+1), '-c:s:%d' % i, 'mov_text'])
476             lang_code = self._conver_lang_code(lang)
477             if lang_code is not None:
478                 opts.extend(['-metadata:s:s:%d' % i, 'language=%s' % lang_code])
479         opts.extend(['-f', 'mp4'])
480
481         temp_filename = filename + u'.temp'
482         self._downloader.to_screen(u'[ffmpeg] Embedding subtitles in \'%s\'' % filename)
483         self.run_ffmpeg_multiple_files(input_files, temp_filename, opts)
484         os.remove(encodeFilename(filename))
485         os.rename(encodeFilename(temp_filename), encodeFilename(filename))
486
487         return True, information
488
489
490 class FFmpegMetadataPP(FFmpegPostProcessor):
491     def run(self, info):
492         metadata = {}
493         if info.get('title') is not None:
494             metadata['title'] = info['title']
495         if info.get('upload_date') is not None:
496             metadata['date'] = info['upload_date']
497         if info.get('uploader') is not None:
498             metadata['artist'] = info['uploader']
499         elif info.get('uploader_id') is not None:
500             metadata['artist'] = info['uploader_id']
501
502         if not metadata:
503             self._downloader.to_screen(u'[ffmpeg] There isn\'t any metadata to add')
504             return True, info
505
506         filename = info['filepath']
507         temp_filename = prepend_extension(filename, 'temp')
508
509         options = ['-c', 'copy']
510         for (name, value) in metadata.items():
511             options.extend(['-metadata', '%s=%s' % (name, value)])
512
513         self._downloader.to_screen(u'[ffmpeg] Adding metadata to \'%s\'' % filename)
514         self.run_ffmpeg(filename, temp_filename, options)
515         os.remove(encodeFilename(filename))
516         os.rename(encodeFilename(temp_filename), encodeFilename(filename))
517         return True, info
518
519
520 class FFmpegMergerPP(FFmpegPostProcessor):
521     def run(self, info):
522         filename = info['filepath']
523         args = ['-c', 'copy']
524         self.run_ffmpeg_multiple_files(info['__files_to_merge'], filename, args)
525         return True, info
526
527
528 class XAttrMetadataPP(PostProcessor):
529
530     #
531     # More info about extended attributes for media:
532     #   http://freedesktop.org/wiki/CommonExtendedAttributes/
533     #   http://www.freedesktop.org/wiki/PhreedomDraft/
534     #   http://dublincore.org/documents/usageguide/elements.shtml
535     #
536     # TODO:
537     #  * capture youtube keywords and put them in 'user.dublincore.subject' (comma-separated)
538     #  * figure out which xattrs can be used for 'duration', 'thumbnail', 'resolution'
539     #
540
541     def run(self, info):
542         """ Set extended attributes on downloaded file (if xattr support is found). """
543
544         # This mess below finds the best xattr tool for the job and creates a
545         # "write_xattr" function.
546         try:
547             # try the pyxattr module...
548             import xattr
549             def write_xattr(path, key, value):
550                 return xattr.setxattr(path, key, value)
551
552         except ImportError:
553
554             if os.name == 'posix':
555                 def which(bin):
556                     for dir in os.environ["PATH"].split(":"):
557                         path = os.path.join(dir, bin)
558                         if os.path.exists(path):
559                             return path
560
561                 user_has_setfattr = which("setfattr")
562                 user_has_xattr    = which("xattr")
563
564                 if user_has_setfattr or user_has_xattr:
565
566                     def write_xattr(path, key, value):
567                         import errno
568                         potential_errors = {
569                             # setfattr: /tmp/blah: Operation not supported
570                             "Operation not supported": errno.EOPNOTSUPP,
571                             # setfattr: ~/blah: No such file or directory
572                             # xattr: No such file: ~/blah
573                             "No such file": errno.ENOENT,
574                         }
575
576                         if user_has_setfattr:
577                             cmd = ['setfattr', '-n', key, '-v', value, path]
578                         elif user_has_xattr:
579                             cmd = ['xattr', '-w', key, value, path]
580
581                         try:
582                             output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
583                         except subprocess.CalledProcessError as e:
584                             errorstr = e.output.strip().decode()
585                             for potential_errorstr, potential_errno in potential_errors.items():
586                                 if errorstr.find(potential_errorstr) > -1:
587                                     e = OSError(potential_errno, potential_errorstr)
588                                     e.__cause__ = None
589                                     raise e
590                             raise # Reraise unhandled error
591
592                 else:
593                     # On Unix, and can't find pyxattr, setfattr, or xattr.
594                     if sys.platform.startswith('linux'):
595                         self._downloader.report_error("Couldn't find a tool to set the xattrs. Install either the python 'pyxattr' or 'xattr' modules, or the GNU 'attr' package (which contains the 'setfattr' tool).")
596                     elif sys.platform == 'darwin':
597                         self._downloader.report_error("Couldn't find a tool to set the xattrs. Install either the python 'xattr' module, or the 'xattr' binary.")
598             else:
599                 # Write xattrs to NTFS Alternate Data Streams: http://en.wikipedia.org/wiki/NTFS#Alternate_data_streams_.28ADS.29
600                 def write_xattr(path, key, value):
601                     assert(key.find(":") < 0)
602                     assert(path.find(":") < 0)
603                     assert(os.path.exists(path))
604
605                     ads_fn = path + ":" + key
606                     with open(ads_fn, "w") as f:
607                         f.write(value)
608
609         # Write the metadata to the file's xattrs
610         self._downloader.to_screen('[metadata] Writing metadata to file\'s xattrs...')
611
612         filename = info['filepath']
613
614         try:
615             xattr_mapping = {
616                 'user.xdg.referrer.url': 'webpage_url',
617                 # 'user.xdg.comment':            'description',
618                 'user.dublincore.title': 'title',
619                 'user.dublincore.date': 'upload_date',
620                 'user.dublincore.description': 'description',
621                 'user.dublincore.contributor': 'uploader',
622                 'user.dublincore.format': 'format',
623             }
624
625             for xattrname, infoname in xattr_mapping.items():
626
627                 value = info.get(infoname)
628
629                 if value:
630                     if infoname == "upload_date":
631                         value = hyphenate_date(value)
632
633                     write_xattr(filename, xattrname, value)
634
635             return True, info
636
637         except OSError:
638             self._downloader.report_error("This filesystem doesn't support extended attributes. (You may have to enable them in your /etc/fstab)")
639             return False, info
640