Merge pull request #5 from ingwinlu/rework_quality

rework quality selection, resolve #4
ingwinlu · Feb 8, 2015 · e9ca8c4 · e9ca8c4
2 parents 0a9dbdf + 3937f48
commit e9ca8c4
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 21 deletions.
diff --git a/twitch/__init__.py b/twitch/__init__.py
@@ -66,44 +66,52 @@ def getJson(self, url, headers=None):
             raise TwitchException(TwitchException.JSON_ERROR)
 
 class M3UPlaylist(object):
-    def __init__(self, input):
+    def __init__(self, input, qualityList = None):
+        self.playlist = dict()
+        self.qualityList = qualityList or Keys.QUALITY_LIST_STREAM
+
         def parseQuality(ExtXMediaLine,ExtXStreamInfLine,Url):
+            #find name of current quality, NAME=", 6 chars
             namePosition = ExtXMediaLine.find('NAME')
             if(namePosition==-1):
                 raise TwitchException()
-            qualityName = ''
+            qualityString = ''
             namePosition+=6
             for char in ExtXMediaLine[namePosition:]:
                 if(char=='"'):
                     break
-                qualityName += char
-            return qualityName, Url
+                qualityString += char
+            return qualityString, Url
 
-        self.playlist = dict()
         lines = input.splitlines()
         linesIterator = iter(lines)
         for line in linesIterator:
             if(line.startswith('#EXT-X-MEDIA')):
                 quality, url = parseQuality(line, next(linesIterator), next(linesIterator))
-                self.playlist[quality] = url
+                qualityInt = self.qualityList.index(quality)
+                self.playlist[qualityInt] = url
+        if not self.playlist:
+            #playlist dict is empty
+            raise ValueError('could not find playable urls')
 
     #returns selected quality or best match if not available
-    def getQuality(self, QualityInt):
-        def isInPlaylist(QualityInt):
-            return Keys.QUALITY_LIST_STREAM[QualityInt] in self.playlist
-
-        if(isInPlaylist(QualityInt)):
+    def getQuality(self, selectedQuality):
+        if(selectedQuality in self.playlist.keys()):
             #selected quality is available
-            return self.playlist[Keys.QUALITY_LIST_STREAM[QualityInt]]
+            return self.playlist[selectedQuality]
         else:
-            #not available, start with worst quality and improve
-            #break if better quality is not available
-            bestMatch = len(Keys.QUALITY_LIST_STREAM) - 1
-            for newMatch in range(bestMatch, -1, -1):
-                if(isInPlaylist(newMatch)):
-                    bestMatch = newMatch
+            #not available, calculate differences to available qualities
+            #return lowest difference / lower quality if same distance
+            bestDistance = len(self.qualityList) + 1
+            bestMatch = None
+
+            for quality in sorted(self.playlist, reverse=True):
+                newDistance = abs(selectedQuality - quality)
+                if newDistance < bestDistance:
+                    bestDistance = newDistance
+                    bestMatch = quality
 
-            return self.playlist[Keys.QUALITY_LIST_STREAM[bestMatch]]
+            return self.playlist[bestMatch]
 
     def __str__(self):
         return repr(self.playlist)
@@ -313,8 +321,9 @@ class Keys(object):
     PREVIEW = 'preview'
     TITLE = 'title'
 
-    QUALITY_LIST_STREAM = ['Source', "High", "Medium", "Low", "Mobile"]
-    QUALITY_LIST_VIDEO = ['live', "720p", "480p", "360p", "226p"]
+    QUALITY_LIST_STREAM = ['Source', 'High', 'Medium', 'Low', 'Mobile']
+    QUALITY_LIST_VIDEO = ['live', '720p', '480p', '360p', '226p']
+
 
 class Urls(object):
     '''

diff --git a/twitch/tests/test_m3u.py b/twitch/tests/test_m3u.py
@@ -39,6 +39,21 @@ class TestM3U(unittest.TestCase):
 #EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=164000,VIDEO="mobile"
 http://video16.prg01.hls.twitch.tv/hls18/ongamenet_9656195424_98434331/mobile/index-live.m3u8?token=id=6125541036366455728,bid=9656195424,exp=1401014215,node=video16-1.prg01.hls.justin.tv,nname=video16.prg01,fmt=mobile&sig=cadce653f0b6b663f2a25c4e946788a30f559447"""
 
+    quality_select = """
+#EXTM3U
+#EXT-X-TWITCH-INFO:NODE="video2.prg01",MANIFEST-NODE="video2.prg01",SERVER-TIME="1400918840.88",USER-IP="84.112.27.151",CLUSTER="prg01",MANIFEST-CLUSTER="prg01"
+#EXT-X-TWITCH-RESTRICTED:GROUP-ID="chunked",NAME="Source",RESTRICTION="chansub"
+#EXT-X-MEDIA:TYPE=VIDEO,GROUP-ID="high",NAME="High",AUTOSELECT=YES,DEFAULT=YES
+#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=1760000,VIDEO="high"
+http://video2.prg01.hls.twitch.tv/hls106/riotgamesoceania_9652805392_98328050/high/index-live.m3u8?token=id=7828074928424501897,bid=9652805392,exp=1401005240,node=video2-1.prg01.hls.justin.tv,nname=video2.prg01,fmt=high&sig=7c20fa2263c78892c0f15c818620f0e85557cabf
+#EXT-X-TWITCH-RESTRICTED:GROUP-ID="medium",NAME="Medium",RESTRICTION="chansub"
+#EXT-X-MEDIA:TYPE=VIDEO,GROUP-ID="low",NAME="Low",AUTOSELECT=YES,DEFAULT=YES
+#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=596000,VIDEO="low"
+http://video2.prg01.hls.twitch.tv/hls106/riotgamesoceania_9652805392_98328050/low/index-live.m3u8?token=id=7828074928424501897,bid=9652805392,exp=1401005240,node=video2-1.prg01.hls.justin.tv,nname=video2.prg01,fmt=low&sig=089be1e11a1556877ca575b3eada7a24acc7ea5a
+#EXT-X-MEDIA:TYPE=VIDEO,GROUP-ID="mobile",NAME="Mobile",AUTOSELECT=YES,DEFAULT=YES
+#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=164000,VIDEO="mobile"
+http://video2.prg01.hls.twitch.tv/hls106/riotgamesoceania_9652805392_98328050/mobile/index-live.m3u8?token=id=7828074928424501897,bid=9652805392,exp=1401005240,node=video2-1.prg01.hls.justin.tv,nname=video2.prg01,fmt=mobile&sig=087b4221fdd0653456b55b5b77756ca3cadd0226"""
+
     vod = """
 #EXTM3U
 #EXT-X-MEDIA:TYPE=VIDEO,GROUP-ID="chunked",NAME="Source",AUTOSELECT=YES,DEFAULT=YES
@@ -84,6 +99,15 @@ def test_vod_0(self):
         url = M3UPlaylist(self.vod).getQuality(0)
         self.assertEqual(url, expected)
 
+    def test_empty_playlist(self):
+        with self.assertRaises(ValueError):
+            M3UPlaylist('')
+
+    def test_bestMatch_quality(self):
+        expected = 'http://video2.prg01.hls.twitch.tv/hls106/riotgamesoceania_9652805392_98328050/low/index-live.m3u8?token=id=7828074928424501897,bid=9652805392,exp=1401005240,node=video2-1.prg01.hls.justin.tv,nname=video2.prg01,fmt=low&sig=089be1e11a1556877ca575b3eada7a24acc7ea5a'
+        url = M3UPlaylist(self.quality_select).getQuality(2)
+        self.assertEqual(url, expected)
+
     def suite(self):
         testSuite = unittest.TestSuite()
         testSuite.addTest(unittest.makeSuite(TestResolver))