From f9c127b989028ccae95b8e6659054dd25dd5fb2e Mon Sep 17 00:00:00 2001
From: Ali Fardan <root@ravenholm.info>
Date: Wed, 5 Feb 2020 09:08:11 +0300
Subject: [PATCH] introduce new '--add-music-metadata' options

---
 youtube_dl/__init__.py               |  2 +
 youtube_dl/options.py                |  4 ++
 youtube_dl/postprocessor/__init__.py |  2 +
 youtube_dl/postprocessor/ffmpeg.py   | 55 ++++++++++++++++++++++++++++
 4 files changed, 63 insertions(+)

diff --git a/youtube_dl/__init__.py b/youtube_dl/__init__.py
index 9a659fc65..299359498 100644
--- a/youtube_dl/__init__.py
+++ b/youtube_dl/__init__.py
@@ -276,6 +276,8 @@ def _real_main(argv=None):
     # so metadata can be added here.
     if opts.addmetadata:
         postprocessors.append({'key': 'FFmpegMetadata'})
+    if opts.addmusicmetadata:
+        postprocessors.append({'key': 'FFmpegMusicMetadata'})
     if opts.convertsubtitles:
         postprocessors.append({
             'key': 'FFmpegSubtitlesConvertor',
diff --git a/youtube_dl/options.py b/youtube_dl/options.py
index 1ffabc62b..8d8deaba9 100644
--- a/youtube_dl/options.py
+++ b/youtube_dl/options.py
@@ -818,6 +818,10 @@ def parseOpts(overrideArguments=None):
         '--add-metadata',
         action='store_true', dest='addmetadata', default=False,
         help='Write metadata to the video file')
+    postproc.add_option(
+        '--add-music-metadata',
+        action='store_true', dest='addmusicmetadata', default=False,
+        help='Write music metadata to video file (this works only for songs from music.youtube.com)')
     postproc.add_option(
         '--metadata-from-title',
         metavar='FORMAT', dest='metafromtitle',
diff --git a/youtube_dl/postprocessor/__init__.py b/youtube_dl/postprocessor/__init__.py
index 3ea518399..dd024b9c7 100644
--- a/youtube_dl/postprocessor/__init__.py
+++ b/youtube_dl/postprocessor/__init__.py
@@ -10,6 +10,7 @@ from .ffmpeg import (
     FFmpegFixupM4aPP,
     FFmpegMergerPP,
     FFmpegMetadataPP,
+    FFmpegMusicMetadataPP,
     FFmpegVideoConvertorPP,
     FFmpegSubtitlesConvertorPP,
 )
@@ -32,6 +33,7 @@ __all__ = [
     'FFmpegFixupStretchedPP',
     'FFmpegMergerPP',
     'FFmpegMetadataPP',
+    'FFmpegMusicMetadataPP',
     'FFmpegPostProcessor',
     'FFmpegSubtitlesConvertorPP',
     'FFmpegVideoConvertorPP',
diff --git a/youtube_dl/postprocessor/ffmpeg.py b/youtube_dl/postprocessor/ffmpeg.py
index fd3f921a8..24d3dea28 100644
--- a/youtube_dl/postprocessor/ffmpeg.py
+++ b/youtube_dl/postprocessor/ffmpeg.py
@@ -502,6 +502,61 @@ class FFmpegMetadataPP(FFmpegPostProcessor):
         os.rename(encodeFilename(temp_filename), encodeFilename(filename))
         return [], info
 
+class FFmpegMusicMetadataPP(FFmpegPostProcessor):
+    def run(self, info):
+        metadata = {}
+
+        # Any original music (even singles) provided from music.youtube.com will
+        # contain 'album' JSON field, if not, its just like any other video
+        if info.get('album') is None:
+            self._downloader.to_screen('[ffmpeg] No music metadata found, nothing to add')
+            return [], info
+        metadata['title'] = info['track']
+        metadata['date'] = info['release_year']
+        metadata['artist'] = info['artist']
+        metadata['album'] = info['album']
+        if info.get('playlist_index') is not None:
+            metadata['track'] = '%d/%d' % (info['playlist_index'], info['n_entries'])
+
+        filename = info['filepath']
+        temp_filename = prepend_extension(filename, 'temp')
+        in_filenames = [filename]
+        options = []
+
+        if info['ext'] == 'm4a':
+            options.extend(['-vn', '-acodec', 'copy'])
+        else:
+            options.extend(['-c', 'copy'])
+
+        for (name, value) in metadata.items():
+            options.extend(['-metadata', '%s=%s' % (name, value)])
+
+        chapters = info.get('chapters', [])
+        if chapters:
+            metadata_filename = replace_extension(filename, 'meta')
+            with io.open(metadata_filename, 'wt', encoding='utf-8') as f:
+                def ffmpeg_escape(text):
+                    return re.sub(r'(=|;|#|\\|\n)', r'\\\1', text)
+
+                metadata_file_content = ';FFMETADATA1\n'
+                for chapter in chapters:
+                    metadata_file_content += '[CHAPTER]\nTIMEBASE=1/1000\n'
+                    metadata_file_content += 'START=%d\n' % (chapter['start_time'] * 1000)
+                    metadata_file_content += 'END=%d\n' % (chapter['end_time'] * 1000)
+                    chapter_title = chapter.get('title')
+                    if chapter_title:
+                        metadata_file_content += 'title=%s\n' % ffmpeg_escape(chapter_title)
+                f.write(metadata_file_content)
+                in_filenames.append(metadata_filename)
+                options.extend(['-map_metadata', '1'])
+
+        self._downloader.to_screen('[ffmpeg] Adding metadata to \'%s\'' % filename)
+        self.run_ffmpeg_multiple_files(in_filenames, temp_filename, options)
+        if chapters:
+            os.remove(metadata_filename)
+        os.remove(encodeFilename(filename))
+        os.rename(encodeFilename(temp_filename), encodeFilename(filename))
+        return [], info
 
 class FFmpegMergerPP(FFmpegPostProcessor):
     def run(self, info):