Smoother JSON checks.

3 files changed, 16 insertions, 15 deletions

diff --git a/src/livestreamer/plugins/svtplay.py b/src/livestreamer/plugins/svtplay.py
index 5757ac1..62c7555 100644
--- a/src/livestreamer/plugins/svtplay.py
+++ b/src/livestreamer/plugins/svtplay.py
@@ -3,7 +3,7 @@
 from livestreamer.compat import str
 from livestreamer.plugins import Plugin, PluginError, NoStreamsError, register_plugin
 from livestreamer.stream import RTMPStream
-from livestreamer.utils import urlget, swfverify
+from livestreamer.utils import urlget, swfverify, verifyjson
 
 import json, re
 
@@ -37,11 +37,9 @@ class SVTPlay(Plugin):
         except ValueError as err:
             raise PluginError(("Unable to parse JSON: {0})").format(err))
 
-        if not ("video" in info and "videoReferences" in info["video"]):
-            raise PluginError("Missing 'video' or 'videoReferences' key in JSON")
-
         streams = {}
-        videos = info["video"]["videoReferences"]
+        video = verifyjson(info, "video")
+        videos = verifyjson(video, "videoReferences")
         swfhash, swfsize = swfverify(self.SWFURL)
 
         for video in videos:
diff --git a/src/livestreamer/plugins/youtube.py b/src/livestreamer/plugins/youtube.py
index 88ae6de..804ddbc 100644
--- a/src/livestreamer/plugins/youtube.py
+++ b/src/livestreamer/plugins/youtube.py
@@ -3,7 +3,7 @@
 from livestreamer.compat import str, bytes, parse_qs
 from livestreamer.plugins import Plugin, PluginError, NoStreamsError, register_plugin
 from livestreamer.stream import HTTPStream
-from livestreamer.utils import urlget
+from livestreamer.utils import urlget, verifyjson
 
 import re
 import json
@@ -61,20 +61,18 @@ class Youtube(Plugin):
         if not info:
             raise NoStreamsError(self.url)
 
-        if "args" in info:
-            args = info["args"]
-        else:
-            raise PluginError("JSON data is missing 'args' key")
+        args = verifyjson(info, "args")
 
         if not "live_playback" in args or args["live_playback"] == "0":
             raise NoStreamsError(self.url)
 
-        if not ("url_encoded_fmt_stream_map" in args and "fmt_list" in args):
-            raise PluginError("JSON data is missing 'url_encoded_fmt_stream_map' or 'fmt_list' keys")
-
         streams = {}
-        streammap = self._parse_stream_map(args["url_encoded_fmt_stream_map"])
-        formatmap = self._parse_format_map(args["fmt_list"])
+
+        uestreammap = verifyjson(args, "url_encoded_fmt_stream_map")
+        fmtlist = verifyjson(args, "fmt_list")
+
+        streammap = self._parse_stream_map(uestreammap)
+        formatmap = self._parse_format_map(fmtlist)
 
         for streaminfo in streammap:
             if not "url" in streaminfo:
diff --git a/src/livestreamer/utils.py b/src/livestreamer/utils.py
index 4cef351..177df39 100644
--- a/src/livestreamer/utils.py
+++ b/src/livestreamer/utils.py
@@ -58,3 +58,8 @@ def swfverify(url):
 
     return h.hexdigest(), len(swf)
 
+def verifyjson(json, key):
+    if not key in json:
+        raise PluginError(("Missing '{0}' key in JSON").format(key))
+
+    return json[key]