Redo av codec settings & selections to accomodate webm

Allows for ranked preferences for h264, av1, and vp9 codecs in
settings, along with equal preferences which are tiebroken using
smaller file size.

For each quality, gives av-merge a list of video sources
and audio sources sorted based on preference & file size. It
will pick the first one that the browser supports.

Closes #84

Signed-off-by: Jesús <heckyel@hyperbola.info>

6 files changed, 126 insertions, 95 deletions

diff --git a/youtube/static/js/av-merge.js b/youtube/static/js/av-merge.js
index 6ba7f79..bbe3e35 100644
--- a/youtube/static/js/av-merge.js
+++ b/youtube/static/js/av-merge.js
@@ -19,13 +19,12 @@
 
 // TODO: Call abort to cancel in-progress appends?
 
-function AVMerge(video, srcPair, startTime){
-    this.videoSource = srcPair[0];
-    this.audioSource = srcPair[1];
-    if (this.videoSource.bitrate && this.audioSource.bitrate)
-        this.avRatio = this.audioSource.bitrate/this.videoSource.bitrate;
-    else
-        this.avRatio = 1/10;
+
+
+function AVMerge(video, srcInfo, startTime){
+    this.audioSource = null;
+    this.videoSource = null;
+    this.avRatio = null;
     this.videoStream = null;
     this.audioStream = null;
     this.seeking = false;
@@ -36,30 +35,48 @@ function AVMerge(video, srcPair, startTime){
     this.opened = false;
     this.audioEndOfStreamCalled = false;
     this.videoEndOfStreamCalled = false;
-    this.setup();
-}
-AVMerge.prototype.setup = function() {
     if (!('MediaSource' in window)) {
         reportError('MediaSource not supported.');
         return;
     }
-    var audioSupported = MediaSource.isTypeSupported(
-        this.audioSource['mime_codec']
-    )
-    var videoSupported = MediaSource.isTypeSupported(
-        this.videoSource['mime_codec']
-    )
-    if (!audioSupported)
-        reportError('Unsupported MIME type or codec: ',
-                    this.audioSource['mime_codec']);
-    if (!videoSupported)
-        reportError('Unsupported MIME type or codec: ',
-                    this.videoSource['mime_codec']);
-    if (audioSupported && videoSupported) {
-        this.mediaSource = new MediaSource();
-        this.video.src = URL.createObjectURL(this.mediaSource);
-        this.mediaSource.onsourceopen = this.sourceOpen.bind(this);
+
+    // Find supported video and audio sources
+    for (var src of srcInfo['videos']) {
+        if (MediaSource.isTypeSupported(src['mime_codec'])) {
+            reportDebug('Using video source', src['mime_codec'],
+                        src['quality_string'], 'itag', src['itag']);
+            this.videoSource = src;
+            break;
+        }
     }
+    for (var src of srcInfo['audios']) {
+        if (MediaSource.isTypeSupported(src['mime_codec'])) {
+            reportDebug('Using audio source', src['mime_codec'],
+                        src['quality_string'], 'itag', src['itag']);
+            this.audioSource = src;
+            break;
+        }
+    }
+    if (this.videoSource === null)
+        reportError('No supported video MIME type or codec found: ',
+                    srcInfo['videos'].map(s => s.mime_codec).join(', '));
+    if (this.audioSource === null)
+        reportError('No supported audio MIME type or codec found: ',
+                    srcInfo['audios'].map(s => s.mime_codec).join(', '));
+    if (this.videoSource === null || this.audioSource === null)
+        return;
+
+    if (this.videoSource.bitrate && this.audioSource.bitrate)
+        this.avRatio = this.audioSource.bitrate/this.videoSource.bitrate;
+    else
+        this.avRatio = 1/10;
+
+    this.setup();
+}
+AVMerge.prototype.setup = function() {
+    this.mediaSource = new MediaSource();
+    this.video.src = URL.createObjectURL(this.mediaSource);
+    this.mediaSource.onsourceopen = this.sourceOpen.bind(this);
 }
 
 AVMerge.prototype.sourceOpen = function(_) {
diff --git a/youtube/static/js/plyr-start.js b/youtube/static/js/plyr-start.js
index a72ab9c..82b32d6 100644
--- a/youtube/static/js/plyr-start.js
+++ b/youtube/static/js/plyr-start.js
@@ -20,10 +20,10 @@
     qualityOptions.push(src.quality_string)
   }
   for (var src of data['pair_sources']) {
-    qualityOptions.push(src[0].quality_string)
+    qualityOptions.push(src.quality_string)
   }
   if (data['using_pair_sources'])
-    qualityDefault = data['pair_sources'][data['pair_idx']][0].quality_string;
+    qualityDefault = data['pair_sources'][data['pair_idx']].quality_string;
   else if (data['uni_sources'].length != 0)
     qualityDefault = data['uni_sources'][data['uni_idx']].quality_string;
   else
@@ -108,7 +108,7 @@
           }
         } else {
           for (var i=0; i < data['pair_sources'].length; i++) {
-            if (data['pair_sources'][i][0].quality_string == quality) {
+            if (data['pair_sources'][i].quality_string == quality) {
               changeQuality({'type': 'pair', 'index': i});
               return;
             }
diff --git a/youtube/static/js/watch.js b/youtube/static/js/watch.js
index f7e6285..0954f5b 100644
--- a/youtube/static/js/watch.js
+++ b/youtube/static/js/watch.js
@@ -4,16 +4,15 @@ function changeQuality(selection) {
     var currentVideoTime = video.currentTime;
     var videoPaused = video.paused;
     var videoSpeed = video.playbackRate;
-    var videoSource;
+    var srcInfo;
     if (avMerge)
         avMerge.close();
     if (selection.type == 'uni'){
-        videoSource = data['uni_sources'][selection.index];
-        video.src = videoSource.url;
+        srcInfo = data['uni_sources'][selection.index];
+        video.src = srcInfo.url;
     } else {
-        let srcPair = data['pair_sources'][selection.index];
-        videoSource = srcPair[0];
-        avMerge = new AVMerge(video, srcPair, currentVideoTime);
+        srcInfo = data['pair_sources'][selection.index];
+        avMerge = new AVMerge(video, srcInfo, currentVideoTime);
     }
     video.currentTime = currentVideoTime;
     if (!videoPaused){
@@ -26,7 +25,6 @@ function changeQuality(selection) {
 var avMerge;
 if (data.using_pair_sources) {
     var srcPair = data['pair_sources'][data['pair_idx']];
-    var videoSource = srcPair[0];
     // Do it dynamically rather than as the default in jinja
     // in case javascript is disabled
     avMerge = new AVMerge(video, srcPair, 0);
diff --git a/youtube/templates/watch.html b/youtube/templates/watch.html
index 46a3645..6aa7e42 100644
--- a/youtube/templates/watch.html
+++ b/youtube/templates/watch.html
@@ -87,7 +87,7 @@
                         <option value='{"type": "uni", "index": {{ loop.index0 }}}' {{ 'selected' if loop.index0 == uni_idx and not using_pair_sources else '' }} >{{ src['quality_string'] }}</option>
                     {% endfor %}
                     {% for src_pair in pair_sources %}
-                        <option value='{"type": "pair", "index": {{ loop.index0}}}' {{ 'selected' if loop.index0 == pair_idx and using_pair_sources else '' }} >{{ src_pair[0]['quality_string'] }}, {{ src_pair[1]['quality_string'] }}</option>
+                        <option value='{"type": "pair", "index": {{ loop.index0}}}' {{ 'selected' if loop.index0 == pair_idx and using_pair_sources else '' }} >{{ src_pair['quality_string'] }}</option>
                     {% endfor %}
                 </select>
             </div>
diff --git a/youtube/watch.py b/youtube/watch.py
index 80885f9..9fd0333 100644
--- a/youtube/watch.py
+++ b/youtube/watch.py
@@ -23,22 +23,31 @@ except FileNotFoundError:
     decrypt_cache = {}
 
 
+def codec_name(vcodec):
+    if vcodec.startswith('avc'):
+        return 'h264'
+    elif vcodec.startswith('av01'):
+        return 'av1'
+    elif vcodec.startswith('vp'):
+        return 'vp'
+    else:
+        return 'unknown'
+
+
 def get_video_sources(info, target_resolution):
     '''return dict with organized sources: {
         'uni_sources': [{}, ...],   # video and audio in one file
         'uni_idx': int,     # default unified source index
-        'pair_sources': [({video}, {audio}), ...],
+        'pair_sources': [{video: {}, audio: {}, quality: ..., ...}, ...],
         'pair_idx': int,    # default pair source index
     }
     '''
     audio_sources = []
-    video_only_sources = []
+    video_only_sources = {}
     uni_sources = []
     pair_sources = []
     for fmt in info['formats']:
-        if not all(fmt[attr] for attr in ('ext', 'url')):
-            continue
-        if fmt['ext'] != 'mp4': # temporary until webm support
+        if not all(fmt[attr] for attr in ('ext', 'url', 'itag')):
             continue
 
         # unified source
@@ -69,8 +78,11 @@ def get_video_sources(info, target_resolution):
             source['mime_codec'] = (source['type'] + '; codecs="'
                                     + source['acodec'] + '"')
             audio_sources.append(source)
-        # video-only source, include audio source
-        elif all(fmt[attr] for attr in ('vcodec', 'quality', 'width')):
+        # video-only source
+        elif all(fmt[attr] for attr in ('vcodec', 'quality', 'width', 'fps',
+                                        'file_size')):
+            if codec_name(fmt['vcodec']) == 'unknown':
+                continue
             source = {
                 'type': 'video/' + fmt['ext'],
                 'quality_string': short_video_quality_string(fmt),
@@ -78,66 +90,61 @@ def get_video_sources(info, target_resolution):
             source.update(fmt)
             source['mime_codec'] = (source['type'] + '; codecs="'
                                     + source['vcodec'] + '"')
-            video_only_sources.append(source)
-
-    # Remove alternative mp4 codecs from video sources
-    def codec_name(vcodec):
-        if vcodec.startswith('avc'):
-            return 'h.264'
-        elif vcodec.startswith('av01'):
-            return 'av1'
-        else:
-            return 'unknown'
-    quality_to_codecs = {}
-    for src in video_only_sources:
-        if src['quality'] in quality_to_codecs:
-            quality_to_codecs[src['quality']].add(codec_name(src['vcodec']))
-        else:
-            quality_to_codecs[src['quality']] = {codec_name(src['vcodec'])}
-    i = 0
-    while i < len(video_only_sources):
-        src = video_only_sources[i]
-        codecs_for_quality = quality_to_codecs[src['quality']]
-        have_both = ('h.264' in codecs_for_quality
-                     and 'av1' in codecs_for_quality)
-        have_one = ('h.264' in codecs_for_quality
-                    or 'av1' in codecs_for_quality)
-        name = codec_name(src['vcodec'])
-        if name == 'unknown' and have_one:
-            del video_only_sources[i]
-            continue
-        if not have_both:
-            i += 1
-            continue
-        if name == 'av1' and settings.preferred_video_codec == 0:
-            del video_only_sources[i]
-        elif name == 'h.264' and settings.preferred_video_codec == 1:
-            del video_only_sources[i]
-        else:
-            i += 1
+            quality = str(fmt['quality']) + 'p' + str(fmt['fps'])
+            if quality in video_only_sources:
+                video_only_sources[quality].append(source)
+            else:
+                video_only_sources[quality] = [source]
 
     audio_sources.sort(key=lambda source: source['audio_bitrate'])
-    video_only_sources.sort(key=lambda src: src['quality'])
     uni_sources.sort(key=lambda src: src['quality'])
 
-    for source in video_only_sources:
+    webm_audios = [a for a in audio_sources if a['ext'] == 'webm']
+    mp4_audios = [a for a in audio_sources if a['ext'] == 'mp4']
+
+    for quality_string, sources in video_only_sources.items():
         # choose an audio source to go with it
         # 0.5 is semiarbitrary empirical constant to spread audio sources
         # between 144p and 1080p. Use something better eventually.
-        target_audio_bitrate = source['quality']*source.get('fps', 30)/30*0.5
-        compat_audios = [a for a in audio_sources if a['ext'] == source['ext']]
-        if compat_audios:
-            closest_audio_source = compat_audios[0]
-            best_err = target_audio_bitrate - compat_audios[0]['audio_bitrate']
+        quality, fps = map(int, quality_string.split('p'))
+        target_audio_bitrate = quality*fps/30*0.5
+        pair_info = {
+            'quality_string': quality_string,
+            'quality': quality,
+            'height': sources[0]['height'],
+            'width': sources[0]['width'],
+            'fps': fps,
+            'videos': sources,
+            'audios': [],
+        }
+        for audio_choices in (webm_audios, mp4_audios):
+            if not audio_choices:
+                continue
+            closest_audio_source = audio_choices[0]
+            best_err = target_audio_bitrate - audio_choices[0]['audio_bitrate']
             best_err = abs(best_err)
-            for audio_source in compat_audios[1:]:
+            for audio_source in audio_choices[1:]:
                 err = abs(audio_source['audio_bitrate'] - target_audio_bitrate)
                 # once err gets worse we have passed the closest one
                 if err > best_err:
                     break
                 best_err = err
                 closest_audio_source = audio_source
-            pair_sources.append((source, closest_audio_source))
+            pair_info['audios'].append(closest_audio_source)
+
+        if not pair_info['audios']:
+            continue
+
+        def video_rank(src):
+            ''' Sort by settings preference. Use file size as tiebreaker '''
+            setting_name = 'codec_rank_' + codec_name(src['vcodec'])
+            return (settings.current_settings_dict[setting_name],
+                    src['file_size'])
+        pair_info['videos'].sort(key=video_rank)
+
+        pair_sources.append(pair_info)
+
+    pair_sources.sort(key=lambda src: src['quality'])
 
     uni_idx = 0 if uni_sources else None
     for i, source in enumerate(uni_sources):
@@ -146,8 +153,8 @@ def get_video_sources(info, target_resolution):
         uni_idx = i
 
     pair_idx = 0 if pair_sources else None
-    for i, source_pair in enumerate(pair_sources):
-        if source_pair[0]['quality'] > target_resolution:
+    for i, pair_info in enumerate(pair_sources):
+        if pair_info['quality'] > target_resolution:
             break
         pair_idx = i
 
@@ -619,8 +626,7 @@ def get_watch_page(video_id=None):
                                            uni_idx, 'width',
                                            default=640)
 
-    pair_quality = yt_data_extract.deep_get(pair_sources, pair_idx, 0,
-                                            'quality')
+    pair_quality = yt_data_extract.deep_get(pair_sources, pair_idx, 'quality')
     uni_quality = yt_data_extract.deep_get(uni_sources, uni_idx, 'quality')
     pair_error = abs((pair_quality or 360) - target_resolution)
     uni_error = abs((uni_quality or 360) - target_resolution)
@@ -634,6 +640,16 @@ def get_watch_page(video_id=None):
     using_pair_sources = (
         bool(pair_sources) and (not uni_sources or closer_to_target == 'pair')
     )
+    if using_pair_sources:
+        video_height = pair_sources[pair_idx]['height']
+        video_width = pair_sources[pair_idx]['width']
+    else:
+        video_height = yt_data_extract.deep_get(
+            uni_sources, uni_idx, 'height', default=360
+        )
+        video_width = yt_data_extract.deep_get(
+            uni_sources, uni_idx, 'width', default=640
+        )
 
     # 1 second per pixel, or the actual video width
     theater_video_target_width = max(640, info['duration'] or 0, video_width)
diff --git a/youtube/yt_data_extract/watch_extraction.py b/youtube/yt_data_extract/watch_extraction.py
index e0af28e..c73033c 100644
--- a/youtube/yt_data_extract/watch_extraction.py
+++ b/youtube/yt_data_extract/watch_extraction.py
@@ -420,7 +420,7 @@ def _extract_formats(info, player_response):
         fmt['vcodec'] = None
         fmt['width'] = yt_fmt.get('width')
         fmt['height'] = yt_fmt.get('height')
-        fmt['file_size'] = yt_fmt.get('contentLength')
+        fmt['file_size'] = extract_int(yt_fmt.get('contentLength'))
         fmt['audio_sample_rate'] = extract_int(yt_fmt.get('audioSampleRate'))
         fmt['duration_ms'] = yt_fmt.get('approxDurationMs')
         fmt['fps'] = yt_fmt.get('fps')