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