Merge pull request streamlink#675 from beardypig/hls-name-fmt

gravyboat · web-flow · commit ff24c65bdbd3 · 2017-03-07T09:44:26.000-08:00
stream.hls: format string name input for parse_variant_playlist
diff --git a/src/streamlink/plugin/plugin.py b/src/streamlink/plugin/plugin.py
@@ -7,6 +7,10 @@
 from ..exceptions import PluginError, NoStreamsError
 from ..options import Options
 
+# FIXME: This is a crude attempt at making a bitrate's
+# weight end up similar to the weight of a resolution.
+# Someone who knows math, please fix.
+BIT_RATE_WEIGHT_RATIO = 2.8
 
 QUALITY_WEIGTHS_EXTRA = {
     "other": {
@@ -37,28 +41,28 @@ def stream_weight(stream):
         if stream in weights:
             return weights[stream], group
 
-    match = re.match("^(\d+)([k]|[p])?(\d*)([\+])?$", stream)
+    match = re.match(r"^(\d+)(k|p)?(\d+)?(\+)?(?:_(\d+)k)?$", stream)
 
     if match:
-        if match.group(2) == "k":
+        name_type = match.group(2)
+        if name_type == "k":  # bit rate
             bitrate = int(match.group(1))
-
-            # FIXME: This is a crude attempt at making a bitrate's
-            # weight end up similar to the weight of a resolution.
-            # Someone who knows math, please fix.
-            weight = bitrate / 2.8
+            weight = bitrate / BIT_RATE_WEIGHT_RATIO
 
             return weight, "bitrate"
 
-        elif match.group(2) == "p":
+        elif name_type == "p":  # resolution
             weight = int(match.group(1))
 
-            if match.group(3):
+            if match.group(3):  # fps eg. 60p or 50p
                 weight += int(match.group(3))
 
             if match.group(4) == "+":
                 weight += 1
 
+            if match.group(5):  # bit rate classifier for resolution
+                weight += int(match.group(5)) / BIT_RATE_WEIGHT_RATIO
+
             return weight, "pixels"
 
     return 0, "none"
@@ -85,7 +89,7 @@ def stream_type_priority(stream_types, stream):
 
 
 def stream_sorting_filter(expr, stream_weight):
-    match = re.match(r"(?P<op><=|>=|<|>)?(?P<value>[\w\+]+)", expr)
+    match = re.match(r"(?P<op><=|>=|<|>)?(?P<value>[\w+]+)", expr)
 
     if not match:
         raise PluginError("Invalid filter expression: {0}".format(expr))
diff --git a/src/streamlink/stream/hls.py b/src/streamlink/stream/hls.py
@@ -271,16 +271,18 @@ def open(self):
     @classmethod
     def parse_variant_playlist(cls, session_, url, name_key="name",
                                name_prefix="", check_streams=False,
-                               force_restart=False,
+                               force_restart=False, name_fmt=None,
                                **request_params):
         """Attempts to parse a variant playlist and return its streams.
 
         :param url: The URL of the variant playlist.
         :param name_key: Prefer to use this key as stream name, valid keys are:
                          name, pixels, bitrate.
         :param name_prefix: Add this prefix to the stream names.
+        :param check_streams: Only allow streams that are accessible.
         :param force_restart: Start at the first segment even for a live stream
-        :param check_streams: Only allow streams that are accesible.
+        :param name_fmt: A format string for the name, allowed format keys are
+                         name, pixels, bitrate.
         """
         logger = session_.logger.new_module("hls.parse_variant_playlist")
         locale = session_.localization
@@ -338,8 +340,11 @@ def parse_variant_playlist(cls, session_, url, name_key="name",
                 else:
                     names["bitrate"] = "{0}k".format(bw / 1000.0)
 
-            stream_name = (names.get(name_key) or names.get("name") or
-                           names.get("pixels") or names.get("bitrate"))
+            if name_fmt:
+                stream_name = name_fmt.format(**names)
+            else:
+                stream_name = (names.get(name_key) or names.get("name") or
+                               names.get("pixels") or names.get("bitrate"))
 
             if not stream_name:
                 continue
diff --git a/tests/test_plugin_stream.py b/tests/test_plugin_stream.py
@@ -2,6 +2,7 @@
 
 from streamlink import Streamlink
 from streamlink.plugins.stream import StreamURL
+from streamlink.plugin.plugin import stream_weight
 from streamlink.stream import *
 
 
@@ -107,5 +108,31 @@ def test_parse_params(self):
             StreamURL._parse_params(""""conn=['B:1', 'S:authMe', 'O:1', 'NN:code:1.23', 'NS:flag:ok', 'O:0']""")
         )
 
+    def test_stream_weight(self):
+        self.assertEqual(
+            (720, "pixels"),
+            stream_weight("720p"))
+        self.assertEqual(
+            (721, "pixels"),
+            stream_weight("720p+"))
+        self.assertEqual(
+            (780, "pixels"),
+            stream_weight("720p60"))
+
+        self.assertTrue(
+            stream_weight("720p+") > stream_weight("720p"))
+        self.assertTrue(
+            stream_weight("720p") == stream_weight("720p"))
+        self.assertTrue(
+            stream_weight("720p_3000k") > stream_weight("720p_2500k"))
+        self.assertTrue(
+            stream_weight("720p60_3000k") > stream_weight("720p_3000k"))
+        self.assertTrue(
+            stream_weight("720p_3000k") < stream_weight("720p+_3000k"))
+
+        self.assertTrue(
+            stream_weight("3000k") > stream_weight("2500k"))
+
+
 if __name__ == "__main__":
     unittest.main()
