[Streams] Refactoring (#703)

cogs/streams.py

@@ -1,12 +1,11 @@
-import discord
 from discord.ext import commands
 from .utils.dataIO import dataIO
 from .utils.chat_formatting import escape_mass_mentions
 from .utils import checks
-from __main__ import send_cmd_help
 from collections import defaultdict
 from string import ascii_letters
 from random import choice
+import discord
 import os
 import re
 import aiohttp
@@ -14,10 +13,30 @@ import asyncio
 import logging
 
 
+class StreamsError(Exception):
+    pass
+
+
+class StreamNotFound(StreamsError):
+    pass
+
+
+class APIError(StreamsError):
+    pass
+
+
+class InvalidCredentials(StreamsError):
+    pass
+
+
+class OfflineStream(StreamsError):
+    pass
+
+
 class Streams:
     """Streams
 
-    Twitch, Hitbox and Beam alerts"""
+    Alerts for a variety of streaming services"""
 
     def __init__(self, bot):
         self.bot = bot
@@ -33,15 +52,16 @@ class Streams:
         stream = escape_mass_mentions(stream)
         regex = r'^(https?\:\/\/)?(www\.)?(hitbox\.tv\/)'
         stream = re.sub(regex, '', stream)
-        online = await self.hitbox_online(stream)
+        try:
-        if isinstance(online, discord.Embed):
+            embed = await self.hitbox_online(stream)
-            await self.bot.say(embed=online)
+        except OfflineStream:
-        elif online is False:
             await self.bot.say(stream + " is offline.")
-        elif online is None:
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
+        except APIError:
+            await self.bot.say("Error contacting the API.")
         else:
-            await self.bot.say("Error.")
+            await self.bot.say(embed=embed)
 
     @commands.command(pass_context=True)
     async def twitch(self, ctx, stream: str):
@@ -49,19 +69,20 @@ class Streams:
         stream = escape_mass_mentions(stream)
         regex = r'^(https?\:\/\/)?(www\.)?(twitch\.tv\/)'
         stream = re.sub(regex, '', stream)
-        online = await self.twitch_online(stream)
+        try:
-        if isinstance(online, discord.Embed):
+            embed = await self.twitch_online(stream)
-            await self.bot.say(embed=online)
+        except OfflineStream:
-        elif online is False:
             await self.bot.say(stream + " is offline.")
-        elif online == 404:
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
-        elif online == 400:
+        except APIError:
+            await self.bot.say("Error contacting the API.")
+        except InvalidCredentials:
             await self.bot.say("Owner: Client-ID is invalid or not set. "
                                "See `{}streamset twitchtoken`"
                                "".format(ctx.prefix))
         else:
-            await self.bot.say("Error.")
+            await self.bot.say(embed=embed)
 
     @commands.command()
     async def beam(self, stream: str):
@@ -69,22 +90,23 @@ class Streams:
         stream = escape_mass_mentions(stream)
         regex = r'^(https?\:\/\/)?(www\.)?(beam\.pro\/)'
         stream = re.sub(regex, '', stream)
-        online = await self.beam_online(stream)
+        try:
-        if isinstance(online, discord.Embed):
+            embed = await self.beam_online(stream)
-            await self.bot.say(embed=online)
+        except OfflineStream:
-        elif online is False:
             await self.bot.say(stream + " is offline.")
-        elif online is None:
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
+        except APIError:
+            await self.bot.say("Error contacting the API.")
         else:
-            await self.bot.say("Error.")
+            await self.bot.say(embed=embed)
 
     @commands.group(pass_context=True, no_pm=True)
     @checks.mod_or_permissions(manage_server=True)
     async def streamalert(self, ctx):
         """Adds/removes stream alerts from the current channel"""
         if ctx.invoked_subcommand is None:
-            await send_cmd_help(ctx)
+            await self.bot.send_cmd_help(ctx)
 
     @streamalert.command(name="twitch", pass_context=True)
     async def twitch_alert(self, ctx, stream: str):
@@ -93,47 +115,31 @@ class Streams:
         regex = r'^(https?\:\/\/)?(www\.)?(twitch\.tv\/)'
         stream = re.sub(regex, '', stream)
         channel = ctx.message.channel
-        check = await self.twitch_online(stream)
+        try:
-        if check == 404:
+            await self.twitch_online(stream)
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
             return
-        elif check == 400:
+        except APIError:
+            await self.bot.say("Error contacting the API.")
+            return
+        except InvalidCredentials:
             await self.bot.say("Owner: Client-ID is invalid or not set. "
                                "See `{}streamset twitchtoken`"
                                "".format(ctx.prefix))
             return
-        elif check == "error":
+        except OfflineStream:
-            await self.bot.say("Couldn't contact Twitch API. Try again later.")
+            pass
-            return
 
-        done = False
+        enabled = self.enable_or_disable_if_active(self.twitch_streams,
+                                                   stream,
+                                                   channel)
 
-        for i, s in enumerate(self.twitch_streams):
+        if enabled:
-            if s["NAME"] == stream:
-                if channel.id in s["CHANNELS"]:
-                    if len(s["CHANNELS"]) == 1:
-                        self.twitch_streams.remove(s)
-                        await self.bot.say("Alert has been removed "
-                                           "from this channel.")
-                        done = True
-                    else:
-                        self.twitch_streams[i]["CHANNELS"].remove(channel.id)
-                        await self.bot.say("Alert has been removed "
-                                           "from this channel.")
-                        done = True
-                else:
-                    self.twitch_streams[i]["CHANNELS"].append(channel.id)
-                    await self.bot.say("Alert activated. I will notify this " +
-                                       "channel everytime {}".format(stream) +
-                                       " is live.")
-                    done = True
-
-        if not done:
-            self.twitch_streams.append(
-                {"CHANNELS": [channel.id],
-                 "NAME": stream, "ALREADY_ONLINE": False})
             await self.bot.say("Alert activated. I will notify this channel "
-                               "everytime {} is live.".format(stream))
+                               "when {} is live.".format(stream))
+        else:
+            await self.bot.say("Alert has been removed from this channel.")
 
         dataIO.save_json("data/streams/twitch.json", self.twitch_streams)
 
@@ -144,42 +150,26 @@ class Streams:
         regex = r'^(https?\:\/\/)?(www\.)?(hitbox\.tv\/)'
         stream = re.sub(regex, '', stream)
         channel = ctx.message.channel
-        check = await self.hitbox_online(stream)
+        try:
-        if check is None:
+            await self.hitbox_online(stream)
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
             return
-        elif check == "error":
+        except APIError:
-            await self.bot.say("Error.")
+            await self.bot.say("Error contacting the API.")
             return
+        except OfflineStream:
+            pass
 
-        done = False
+        enabled = self.enable_or_disable_if_active(self.hitbox_streams,
+                                                   stream,
+                                                   channel)
 
-        for i, s in enumerate(self.hitbox_streams):
+        if enabled:
-            if s["NAME"] == stream:
-                if channel.id in s["CHANNELS"]:
-                    if len(s["CHANNELS"]) == 1:
-                        self.hitbox_streams.remove(s)
-                        await self.bot.say("Alert has been removed from this "
-                                           "channel.")
-                        done = True
-                    else:
-                        self.hitbox_streams[i]["CHANNELS"].remove(channel.id)
-                        await self.bot.say("Alert has been removed from this "
-                                           "channel.")
-                        done = True
-                else:
-                    self.hitbox_streams[i]["CHANNELS"].append(channel.id)
-                    await self.bot.say("Alert activated. I will notify this "
-                                       "channel everytime "
-                                       "{} is live.".format(stream))
-                    done = True
-
-        if not done:
-            self.hitbox_streams.append(
-                {"CHANNELS": [channel.id], "NAME": stream,
-                 "ALREADY_ONLINE": False})
             await self.bot.say("Alert activated. I will notify this channel "
-                               "everytime {} is live.".format(stream))
+                               "when {} is live.".format(stream))
+        else:
+            await self.bot.say("Alert has been removed from this channel.")
 
         dataIO.save_json("data/streams/hitbox.json", self.hitbox_streams)
 
@@ -190,42 +180,26 @@ class Streams:
         regex = r'^(https?\:\/\/)?(www\.)?(beam\.pro\/)'
         stream = re.sub(regex, '', stream)
         channel = ctx.message.channel
-        check = await self.beam_online(stream)
+        try:
-        if check is None:
+            await self.beam_online(stream)
+        except StreamNotFound:
             await self.bot.say("That stream doesn't exist.")
             return
-        elif check == "error":
+        except APIError:
-            await self.bot.say("Error.")
+            await self.bot.say("Error contacting the API.")
             return
+        except OfflineStream:
+            pass
 
-        done = False
+        enabled = self.enable_or_disable_if_active(self.beam_streams,
+                                                   stream,
+                                                   channel)
 
-        for i, s in enumerate(self.beam_streams):
+        if enabled:
-            if s["NAME"] == stream:
-                if channel.id in s["CHANNELS"]:
-                    if len(s["CHANNELS"]) == 1:
-                        self.beam_streams.remove(s)
-                        await self.bot.say("Alert has been removed from this "
-                                           "channel.")
-                        done = True
-                    else:
-                        self.beam_streams[i]["CHANNELS"].remove(channel.id)
-                        await self.bot.say("Alert has been removed from this "
-                                           "channel.")
-                        done = True
-                else:
-                    self.beam_streams[i]["CHANNELS"].append(channel.id)
-                    await self.bot.say("Alert activated. I will notify this "
-                                       "channel everytime "
-                                       "{} is live.".format(stream))
-                    done = True
-
-        if not done:
-            self.beam_streams.append(
-                {"CHANNELS": [channel.id], "NAME": stream,
-                 "ALREADY_ONLINE": False})
             await self.bot.say("Alert activated. I will notify this channel "
-                               "everytime {} is live.".format(stream))
+                               "when {} is live.".format(stream))
+        else:
+            await self.bot.say("Alert has been removed from this channel.")
 
         dataIO.save_json("data/streams/beam.json", self.beam_streams)
 
@@ -234,41 +208,23 @@ class Streams:
         """Stops all streams alerts in the current channel"""
         channel = ctx.message.channel
 
+        streams = (
+            self.hitbox_streams,
+            self.twitch_streams,
+            self.beam_streams
+        )
+
+        for stream_type in streams:
             to_delete = []
 
-        for s in self.hitbox_streams:
+            for s in stream_type:
                 if channel.id in s["CHANNELS"]:
-                if len(s["CHANNELS"]) == 1:
-                    to_delete.append(s)
-                else:
                     s["CHANNELS"].remove(channel.id)
+                    if not s["CHANNELS"]:
+                        to_delete.append(s)
 
             for s in to_delete:
-            self.hitbox_streams.remove(s)
+                stream_type.remove(s)
-
-        to_delete = []
-
-        for s in self.twitch_streams:
-            if channel.id in s["CHANNELS"]:
-                if len(s["CHANNELS"]) == 1:
-                    to_delete.append(s)
-                else:
-                    s["CHANNELS"].remove(channel.id)
-
-        for s in to_delete:
-            self.twitch_streams.remove(s)
-
-        to_delete = []
-
-        for s in self.beam_streams:
-            if channel.id in s["CHANNELS"]:
-                if len(s["CHANNELS"]) == 1:
-                    to_delete.append(s)
-                else:
-                    s["CHANNELS"].remove(channel.id)
-
-        for s in to_delete:
-            self.beam_streams.remove(s)
 
         dataIO.save_json("data/streams/twitch.json", self.twitch_streams)
         dataIO.save_json("data/streams/hitbox.json", self.hitbox_streams)
@@ -281,7 +237,7 @@ class Streams:
     async def streamset(self, ctx):
         """Stream settings"""
         if ctx.invoked_subcommand is None:
-            await send_cmd_help(ctx)
+            await self.bot.send_cmd_help(ctx)
 
     @streamset.command()
     @checks.is_owner()
@@ -316,57 +272,55 @@ class Streams:
 
     async def hitbox_online(self, stream):
         url = "https://api.hitbox.tv/media/live/" + stream
-        try:
+
         async with aiohttp.get(url) as r:
             data = await r.json(encoding='utf-8')
+
         if "livestream" not in data:
-                return None
+            raise StreamNotFound()
-            if data["livestream"][0]["media_is_live"] == "0":
+        elif data["livestream"][0]["media_is_live"] == "0":
-                return False
+            raise OfflineStream()
         elif data["livestream"][0]["media_is_live"] == "1":
             data = self.hitbox_embed(data)
             return data
-            return "error"
+
-        except:
+        raise APIError()
-            return "error"
 
     async def twitch_online(self, stream):
         session = aiohttp.ClientSession()
         url = "https://api.twitch.tv/kraken/streams/" + stream
         header = {'Client-ID': self.settings.get("TWITCH_TOKEN", "")}
-        try:
+
         async with session.get(url, headers=header) as r:
             data = await r.json(encoding='utf-8')
         await session.close()
-            if r.status == 400:
+        if r.status == 200:
-                return 400
+            if data["stream"] is None:
-            elif r.status == 404:
+                raise OfflineStream()
-                return 404
-            elif data["stream"] is None:
-                return False
-            elif data["stream"]:
             embed = self.twitch_embed(data)
             return embed
-        except:
+        elif r.status == 400:
-            return "error"
+            raise InvalidCredentials()
-        return "error"
+        elif r.status == 404:
+            raise StreamNotFound()
+        else:
+            raise APIError()
 
     async def beam_online(self, stream):
         url = "https://beam.pro/api/v1/channels/" + stream
-        try:
+
         async with aiohttp.get(url) as r:
             data = await r.json(encoding='utf-8')
-            if "online" in data:
+        if r.status == 200:
             if data["online"] is True:
                 data = self.beam_embed(data)
                 return data
             else:
-                    return False
+                raise OfflineStream()
-            elif "error" in data:
+        elif r.status == 404:
-                return None
+            raise StreamNotFound()
-        except:
+        else:
-            return "error"
+            raise APIError()
-        return "error"
 
     def twitch_embed(self, data):
         channel = data["stream"]["channel"]
@@ -425,6 +379,27 @@ class Streams:
             embed.set_footer(text="Playing: " + data["type"]["name"])
         return embed
 
+    def enable_or_disable_if_active(self, streams, stream, channel):
+        """Returns True if enabled or False if disabled"""
+        for i, s in enumerate(streams):
+            if s["NAME"] != stream:
+                continue
+
+            if channel.id in s["CHANNELS"]:
+                streams[i]["CHANNELS"].remove(channel.id)
+                if not s["CHANNELS"]:
+                    streams.remove(s)
+                return False
+            else:
+                streams[i]["CHANNELS"].append(channel.id)
+                return True
+
+        streams.append({"CHANNELS": [channel.id],
+                        "NAME": stream,
+                        "ALREADY_ONLINE": False})
+
+        return True
+
     async def stream_checker(self):
         CHECK_DELAY = 60
 
@@ -435,26 +410,30 @@ class Streams:
                        (self.hitbox_streams, self.hitbox_online),
                        (self.beam_streams,   self.beam_online))
 
-            for stream_type in streams:
+            for streams_list, parser in streams:
-                streams = stream_type[0]
+                for stream in streams_list:
-                parser = stream_type[1]
+                    try:
-                for stream in streams:
+                        embed = await parser(stream["NAME"])
-                    online = await parser(stream["NAME"])
+                    except OfflineStream:
-                    if isinstance(online, discord.Embed) and not stream["ALREADY_ONLINE"]:
+                        if stream["ALREADY_ONLINE"]:
+                            stream["ALREADY_ONLINE"] = False
+                            save = True
+                    except:  # We don't want our task to die
+                        continue
+                    else:
+                        if stream["ALREADY_ONLINE"]:
+                            continue
                         save = True
                         stream["ALREADY_ONLINE"] = True
-                        for channel in stream["CHANNELS"]:
+                        for channel_id in stream["CHANNELS"]:
-                            channel_obj = self.bot.get_channel(channel)
+                            channel = self.bot.get_channel(channel_id)
-                            if channel_obj is None:
+                            if channel is None:
                                 continue
-                            mention = self.settings.get(channel_obj.server.id, {}).get("MENTION", "")
+                            mention = self.settings.get(channel.server.id, {}).get("MENTION", "")
-                            can_speak = channel_obj.permissions_for(channel_obj.server.me).send_messages
+                            can_speak = channel.permissions_for(channel.server.me).send_messages
-                            if channel_obj and can_speak:
+                            if channel and can_speak:
-                                await self.bot.send_message(channel_obj, mention, embed=online)
+                                await self.bot.send_message(channel, mention, embed=embed)
-                    else:
+
-                        if stream["ALREADY_ONLINE"] and not online:
-                            save = True
-                            stream["ALREADY_ONLINE"] = False
                     await asyncio.sleep(0.5)
 
             if save:
@@ -476,20 +455,16 @@ def check_folders():
 
 
 def check_files():
-    f = "data/streams/twitch.json"
+    stream_files = (
-    if not dataIO.is_valid_json(f):
+        "twitch.json",
-        print("Creating empty twitch.json...")
+        "hitbox.json",
-        dataIO.save_json(f, [])
+        "beam.json"
+    )
 
-    f = "data/streams/hitbox.json"
+    for filename in stream_files:
-    if not dataIO.is_valid_json(f):
+        if not dataIO.is_valid_json("data/streams/" + filename):
-        print("Creating empty hitbox.json...")
+            print("Creating empty {}...".format(filename))
-        dataIO.save_json(f, [])
+            dataIO.save_json("data/streams/" + filename, [])
-
-    f = "data/streams/beam.json"
-    if not dataIO.is_valid_json(f):
-        print("Creating empty beam.json...")
-        dataIO.save_json(f, [])
 
     f = "data/streams/settings.json"
     if not dataIO.is_valid_json(f):