From 7a85dc7d74c00f96aedd06a18b3391b1ff343081 Mon Sep 17 00:00:00 2001
From: Johan Andersson <j@i19.se>
Date: Sun, 18 Apr 2021 14:06:25 +0200
Subject: [PATCH] Add resolution info to --list-quality

---
 lib/svtplay_dl/fetcher/__init__.py |  1 +
 lib/svtplay_dl/fetcher/dash.py     | 17 ++++++++++++++++-
 lib/svtplay_dl/fetcher/hls.py      |  4 ++++
 lib/svtplay_dl/utils/stream.py     |  6 +++---
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/lib/svtplay_dl/fetcher/__init__.py b/lib/svtplay_dl/fetcher/__init__.py
index b7bde85..b014bc8 100644
--- a/lib/svtplay_dl/fetcher/__init__.py
+++ b/lib/svtplay_dl/fetcher/__init__.py
@@ -23,6 +23,7 @@ class VideoRetriever:
         self.output_extention = None
         channels = kwargs.pop("channels", None)
         codec = kwargs.pop("codec", "h264")
+        self.resolution = kwargs.pop("resolution", "")
         self.format = f"{codec}-{channels}" if channels else codec
 
     def __repr__(self):
diff --git a/lib/svtplay_dl/fetcher/dash.py b/lib/svtplay_dl/fetcher/dash.py
index c885b19..3b3bac1 100644
--- a/lib/svtplay_dl/fetcher/dash.py
+++ b/lib/svtplay_dl/fetcher/dash.py
@@ -130,6 +130,10 @@ def adaptionset(attributes, elements, url, baseurl=None):
         if "lang" in element.attrib:
             lang = element.attrib["lang"]
 
+        resolution = None
+        if "maxWidth" in element.attrib and "maxHeight" in element.attrib:
+            resolution = f'{element.attrib["maxWidth"]}x{element.attrib["maxHeight"]}'
+
         for i in represtation:
             files = []
             segments = False
@@ -154,6 +158,8 @@ def adaptionset(attributes, elements, url, baseurl=None):
                 codec = "hevc"
             else:
                 codec = codecs
+            if not resolution and "maxWidth" in i.attrib and "maxHeight" in i.attrib:
+                resolution = f'{element.attrib["maxWidth"]}x{element.attrib["maxHeight"]}'
             if i.find("{urn:mpeg:dash:schema:mpd:2011}AudioChannelConfiguration") is not None:
                 chan = i.find("{urn:mpeg:dash:schema:mpd:2011}AudioChannelConfiguration").attrib["value"]
                 if chan == "6":
@@ -176,7 +182,15 @@ def adaptionset(attributes, elements, url, baseurl=None):
                 files.append(filename)
 
             if files:
-                streams[bitrate] = {"segments": segments, "files": files, "codecs": codec, "channels": channels, "lang": lang, "mimetype": mimetype}
+                streams[bitrate] = {
+                    "segments": segments,
+                    "files": files,
+                    "codecs": codec,
+                    "channels": channels,
+                    "lang": lang,
+                    "mimetype": mimetype,
+                    "resolution": resolution,
+                }
 
     return streams
 
@@ -250,6 +264,7 @@ def _dashparse(config, text, url, cookies, **kwargs):
             segments=videofiles[i]["segments"],
             codec=videofiles[i]["codecs"],
             channels=audiofiles[list(audiofiles.keys())[0]]["channels"],
+            resolution=videofiles[i]["resolution"],
             **kwargs,
         )
     for i in subtitles.keys():
diff --git a/lib/svtplay_dl/fetcher/hls.py b/lib/svtplay_dl/fetcher/hls.py
index 1a73d1e..782ad6b 100644
--- a/lib/svtplay_dl/fetcher/hls.py
+++ b/lib/svtplay_dl/fetcher/hls.py
@@ -61,6 +61,7 @@ def hlsparse(config, res, url, **kwargs):
             audio_url = None
             vcodec = None
             chans = None
+            resolution = ""
             if i["TAG"] == "EXT-X-MEDIA":
                 if "AUTOSELECT" in i and (i["AUTOSELECT"].upper() == "YES"):
                     if i["TYPE"] and i["TYPE"] != "SUBTITLES":
@@ -88,6 +89,8 @@ def hlsparse(config, res, url, **kwargs):
                     bit_rate = float(i["AVERAGE-BANDWIDTH"]) / 1000
                 else:
                     bit_rate = float(i["BANDWIDTH"]) / 1000
+                if "RESOLUTION" in i:
+                    resolution = i["RESOLUTION"]
                 if "CODECS" in i:
                     if i["CODECS"][:3] == "hvc":
                         vcodec = "hevc"
@@ -114,6 +117,7 @@ def hlsparse(config, res, url, **kwargs):
                 segments=bool(segments),
                 channels=chans,
                 codec=codec,
+                resolution=resolution,
                 **kwargs,
             )
 
diff --git a/lib/svtplay_dl/utils/stream.py b/lib/svtplay_dl/utils/stream.py
index b4eda0f..53e65f4 100644
--- a/lib/svtplay_dl/utils/stream.py
+++ b/lib/svtplay_dl/utils/stream.py
@@ -15,15 +15,15 @@ def sort_quality(data):
     data = sorted(data, key=lambda x: (x.bitrate, x.name), reverse=True)
     datas = []
     for i in data:
-        datas.append([i.bitrate, i.name, i.format])
+        datas.append([i.bitrate, i.name, i.format, i.resolution])
     return datas
 
 
 def list_quality(videos):
     data = sort_quality(videos)
-    logging.info("Quality\tMethod")
+    logging.info("Quality\tMethod\tCodec\tResolution")
     for i in data:
-        logging.info("%s\t%s\t%s", i[0], i[1].upper(), i[2].upper())
+        logging.info("%s\t%s\t%s\t%s", i[0], i[1].upper(), i[2].upper(), i[3])
 
 
 def protocol_prio(streams, priolist):