From eedda5252c05327748dede204a8fccafa0288118 Mon Sep 17 00:00:00 2001
From: pukkandan <pukkandan.ytdlp@gmail.com>
Date: Mon, 19 Jun 2023 14:06:39 +0530
Subject: [PATCH] [utils] `FormatSorter`: Improve `size` and `br`

Closes #1596

Previously, when some formats have accurate size and some approximate,
the ones with accurate size was always prioritized

For formats with known tbr and unknown vbr/abr, we were setting
(vbr=tbr, abr=0) for sorting to work. This is no longer needed.

Authored by pukkandan, u-spec-png
---
 test/test_InfoExtractor.py | 10 ----------
 yt_dlp/utils/_utils.py     | 25 +++++++++++++++----------
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/test/test_InfoExtractor.py b/test/test_InfoExtractor.py
index 1f60abfd2..b7dee496a 100644
--- a/test/test_InfoExtractor.py
+++ b/test/test_InfoExtractor.py
@@ -917,8 +917,6 @@ def test_parse_m3u8_formats(self):
                     'acodec': 'mp4a.40.2',
                     'video_ext': 'mp4',
                     'audio_ext': 'none',
-                    'vbr': 263.851,
-                    'abr': 0,
                 }, {
                     'format_id': '577',
                     'format_index': None,
@@ -936,8 +934,6 @@ def test_parse_m3u8_formats(self):
                     'acodec': 'mp4a.40.2',
                     'video_ext': 'mp4',
                     'audio_ext': 'none',
-                    'vbr': 577.61,
-                    'abr': 0,
                 }, {
                     'format_id': '915',
                     'format_index': None,
@@ -955,8 +951,6 @@ def test_parse_m3u8_formats(self):
                     'acodec': 'mp4a.40.2',
                     'video_ext': 'mp4',
                     'audio_ext': 'none',
-                    'vbr': 915.905,
-                    'abr': 0,
                 }, {
                     'format_id': '1030',
                     'format_index': None,
@@ -974,8 +968,6 @@ def test_parse_m3u8_formats(self):
                     'acodec': 'mp4a.40.2',
                     'video_ext': 'mp4',
                     'audio_ext': 'none',
-                    'vbr': 1030.138,
-                    'abr': 0,
                 }, {
                     'format_id': '1924',
                     'format_index': None,
@@ -993,8 +985,6 @@ def test_parse_m3u8_formats(self):
                     'acodec': 'mp4a.40.2',
                     'video_ext': 'mp4',
                     'audio_ext': 'none',
-                    'vbr': 1924.009,
-                    'abr': 0,
                 }],
                 {
                     'en': [{
diff --git a/yt_dlp/utils/_utils.py b/yt_dlp/utils/_utils.py
index 646210116..1fd6f44af 100644
--- a/yt_dlp/utils/_utils.py
+++ b/yt_dlp/utils/_utils.py
@@ -5669,6 +5669,7 @@ def orderedSet_from_options(options, alias_dict, *, use_regex=False, start=None)
     return orderedSet(requested)
 
 
+# TODO: Rewrite
 class FormatSorter:
     regex = r' *((?P<reverse>\+)?(?P<field>[a-zA-Z0-9_]+)((?P<separator>[~:])(?P<limit>.*?))?)? *$'
 
@@ -5717,8 +5718,10 @@ class FormatSorter:
         'source': {'convert': 'float', 'field': 'source_preference', 'default': -1},
 
         'codec': {'type': 'combined', 'field': ('vcodec', 'acodec')},
-        'br': {'type': 'combined', 'field': ('tbr', 'vbr', 'abr'), 'same_limit': True},
-        'size': {'type': 'combined', 'same_limit': True, 'field': ('filesize', 'fs_approx')},
+        'br': {'type': 'multiple', 'field': ('tbr', 'vbr', 'abr'),
+               'function': lambda it: next(filter(None, it), None)},
+        'size': {'type': 'multiple', 'field': ('filesize', 'fs_approx'),
+                 'function': lambda it: next(filter(None, it), None)},
         'ext': {'type': 'combined', 'field': ('vext', 'aext')},
         'res': {'type': 'multiple', 'field': ('height', 'width'),
                 'function': lambda it: (lambda l: min(l) if l else 0)(tuple(filter(None, it)))},
@@ -5949,13 +5952,15 @@ def calculate_preference(self, format):
             format['preference'] = -100
 
         # Determine missing bitrates
-        if format.get('tbr') is None:
-            if format.get('vbr') is not None and format.get('abr') is not None:
-                format['tbr'] = format.get('vbr', 0) + format.get('abr', 0)
-        else:
-            if format.get('vcodec') != 'none' and format.get('vbr') is None:
-                format['vbr'] = format.get('tbr') - format.get('abr', 0)
-            if format.get('acodec') != 'none' and format.get('abr') is None:
-                format['abr'] = format.get('tbr') - format.get('vbr', 0)
+        if format.get('vcodec') == 'none':
+            format['vbr'] = 0
+        if format.get('acodec') == 'none':
+            format['abr'] = 0
+        if not format.get('vbr') and format.get('vcodec') != 'none':
+            format['vbr'] = try_call(lambda: format['tbr'] - format['abr']) or None
+        if not format.get('abr') and format.get('acodec') != 'none':
+            format['abr'] = try_call(lambda: format['tbr'] - format['vbr']) or None
+        if not format.get('tbr'):
+            format['tbr'] = try_call(lambda: format['vbr'] + format['abr']) or None
 
         return tuple(self._calculate_field_preference(format, field) for field in self._order)