From f9d0abef9270c9f318188ffeefd8b821c9dbd044 Mon Sep 17 00:00:00 2001 From: Twentysix Date: Thu, 13 Apr 2017 22:42:14 +0200 Subject: [PATCH] [Streams] Refactoring (#703) --- cogs/streams.py | 417 +++++++++++++++++++++++------------------------- 1 file changed, 196 insertions(+), 221 deletions(-) diff --git a/cogs/streams.py b/cogs/streams.py index 6b4975bb2..618864c74 100644 --- a/cogs/streams.py +++ b/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) - if isinstance(online, discord.Embed): - await self.bot.say(embed=online) - elif online is False: + try: + embed = await self.hitbox_online(stream) + except OfflineStream: 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) - if isinstance(online, discord.Embed): - await self.bot.say(embed=online) - elif online is False: + try: + embed = await self.twitch_online(stream) + except OfflineStream: 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) - if isinstance(online, discord.Embed): - await self.bot.say(embed=online) - elif online is False: + try: + embed = await self.beam_online(stream) + except OfflineStream: 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) - if check == 404: + try: + 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": - await self.bot.say("Couldn't contact Twitch API. Try again later.") - return + except OfflineStream: + pass - done = False + enabled = self.enable_or_disable_if_active(self.twitch_streams, + stream, + channel) - for i, s in enumerate(self.twitch_streams): - 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}) + if enabled: 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) - if check is None: + try: + await self.hitbox_online(stream) + except StreamNotFound: await self.bot.say("That stream doesn't exist.") return - elif check == "error": - await self.bot.say("Error.") + except APIError: + 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 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}) + if enabled: 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) - if check is None: + try: + await self.beam_online(stream) + except StreamNotFound: await self.bot.say("That stream doesn't exist.") return - elif check == "error": - await self.bot.say("Error.") + except APIError: + 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 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}) + if enabled: 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 - to_delete = [] + streams = ( + self.hitbox_streams, + self.twitch_streams, + self.beam_streams + ) - for s in self.hitbox_streams: - if channel.id in s["CHANNELS"]: - if len(s["CHANNELS"]) == 1: - to_delete.append(s) - else: + for stream_type in streams: + to_delete = [] + + for s in stream_type: + if channel.id in s["CHANNELS"]: s["CHANNELS"].remove(channel.id) + if not s["CHANNELS"]: + to_delete.append(s) - for s in to_delete: - self.hitbox_streams.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) + for s in to_delete: + stream_type.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 - if data["livestream"][0]["media_is_live"] == "0": - return False - elif data["livestream"][0]["media_is_live"] == "1": - data = self.hitbox_embed(data) - return data - return "error" - except: - return "error" + + async with aiohttp.get(url) as r: + data = await r.json(encoding='utf-8') + + if "livestream" not in data: + raise StreamNotFound() + elif data["livestream"][0]["media_is_live"] == "0": + raise OfflineStream() + elif data["livestream"][0]["media_is_live"] == "1": + data = self.hitbox_embed(data) + return data + + raise APIError() 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: - return 400 - elif r.status == 404: - return 404 - elif data["stream"] is None: - return False - elif data["stream"]: - embed = self.twitch_embed(data) - return embed - except: - return "error" - return "error" + + async with session.get(url, headers=header) as r: + data = await r.json(encoding='utf-8') + await session.close() + if r.status == 200: + if data["stream"] is None: + raise OfflineStream() + embed = self.twitch_embed(data) + return embed + elif r.status == 400: + raise InvalidCredentials() + 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 data["online"] is True: - data = self.beam_embed(data) - return data - else: - return False - elif "error" in data: - return None - except: - return "error" - return "error" + + async with aiohttp.get(url) as r: + data = await r.json(encoding='utf-8') + if r.status == 200: + if data["online"] is True: + data = self.beam_embed(data) + return data + else: + raise OfflineStream() + elif r.status == 404: + raise StreamNotFound() + else: + raise APIError() 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: - streams = stream_type[0] - parser = stream_type[1] - for stream in streams: - online = await parser(stream["NAME"]) - if isinstance(online, discord.Embed) and not stream["ALREADY_ONLINE"]: + for streams_list, parser in streams: + for stream in streams_list: + try: + embed = await parser(stream["NAME"]) + except OfflineStream: + 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"]: - channel_obj = self.bot.get_channel(channel) - if channel_obj is None: + for channel_id in stream["CHANNELS"]: + channel = self.bot.get_channel(channel_id) + if channel is None: continue - mention = self.settings.get(channel_obj.server.id, {}).get("MENTION", "") - can_speak = channel_obj.permissions_for(channel_obj.server.me).send_messages - if channel_obj and can_speak: - await self.bot.send_message(channel_obj, mention, embed=online) - else: - if stream["ALREADY_ONLINE"] and not online: - save = True - stream["ALREADY_ONLINE"] = False + mention = self.settings.get(channel.server.id, {}).get("MENTION", "") + can_speak = channel.permissions_for(channel.server.me).send_messages + if channel and can_speak: + await self.bot.send_message(channel, mention, embed=embed) + await asyncio.sleep(0.5) if save: @@ -476,20 +455,16 @@ def check_folders(): def check_files(): - f = "data/streams/twitch.json" - if not dataIO.is_valid_json(f): - print("Creating empty twitch.json...") - dataIO.save_json(f, []) + stream_files = ( + "twitch.json", + "hitbox.json", + "beam.json" + ) - f = "data/streams/hitbox.json" - if not dataIO.is_valid_json(f): - print("Creating empty hitbox.json...") - dataIO.save_json(f, []) - - f = "data/streams/beam.json" - if not dataIO.is_valid_json(f): - print("Creating empty beam.json...") - dataIO.save_json(f, []) + for filename in stream_files: + if not dataIO.is_valid_json("data/streams/" + filename): + print("Creating empty {}...".format(filename)) + dataIO.save_json("data/streams/" + filename, []) f = "data/streams/settings.json" if not dataIO.is_valid_json(f):