diff --git a/Main.py b/Main.py index c990993d..94ca4d5d 100755 --- a/Main.py +++ b/Main.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- +#! /usr/bin/env python3 # -# Copyright (C) idoneam (2016-2019) +# Copyright (C) idoneam (2016-2020) # # This file is part of Canary # @@ -167,7 +166,7 @@ async def backup(ctx): bot.dev_logger.info('Database backup') -if __name__ == "__main__": +def main(): for extension in startup: try: bot.load_extension(extension) @@ -175,3 +174,7 @@ async def backup(ctx): bot.dev_logger.warning(f'Failed to load extension {extension}\n' f'{type(e).__name__}: {e}') bot.run(bot.config.discord_key) + + +if __name__ == "__main__": + main() diff --git a/bot.py b/bot.py index f7570ce9..3579f599 100644 --- a/bot.py +++ b/bot.py @@ -117,6 +117,11 @@ def _start_database(self): conn.close() self.dev_logger.debug('Database is ready') + def log_traceback(self, exception): + self.dev_logger.error("".join( + traceback.format_exception(type(exception), exception, + exception.__traceback__))) + async def on_command_error(self, ctx, error): """The event triggered when an error is raised while invoking a command. ctx : Context @@ -149,9 +154,7 @@ async def on_command_error(self, ctx, error): self.dev_logger.error('Ignoring exception in command {}:'.format( ctx.command)) - self.dev_logger.error(''.join( - traceback.format_exception(type(error), error, - error.__traceback__))) + self.log_traceback(error) # predefined variables to be imported diff --git a/cogs/subscribers.py b/cogs/subscribers.py index b6643b00..a079d961 100644 --- a/cogs/subscribers.py +++ b/cogs/subscribers.py @@ -1,6 +1,4 @@ -# -*- coding: utf-8 -*- -# -# Copyright (C) idoneam (2016-2019) +# Copyright (C) idoneam (2016-2020) # # This file is part of Canary # @@ -21,19 +19,24 @@ import discord from discord.ext import commands from discord import utils -import asyncio # URL access and parsing from bs4 import BeautifulSoup # Other utilities +import json.decoder import os import re import pickle import feedparser import requests +# Subscriber decorator +from .utils.subscribers import canary_subscriber + CFIA_FEED_URL = "http://inspection.gc.ca/eng/1388422350443/1388422374046.xml" +CFIA_RECALL_TAG_PATH = "pickles/recall_tag.obj" + METRO_STATUS_API = "https://www.stm.info/en/ajax/etats-du-service" METRO_GREEN_LINE = "1" @@ -52,105 +55,131 @@ METRO_NORMAL_SERVICE_MESSAGE = "Normal métro service" -# Default values by line number for status -metro_status = { - METRO_GREEN_LINE: METRO_NORMAL_SERVICE_MESSAGE, - METRO_ORANGE_LINE: METRO_NORMAL_SERVICE_MESSAGE, - METRO_YELLOW_LINE: METRO_NORMAL_SERVICE_MESSAGE, - METRO_BLUE_LINE: METRO_NORMAL_SERVICE_MESSAGE, -} - -os.makedirs('./pickles', exist_ok=True) +os.makedirs("./pickles", exist_ok=True) class Subscribers(commands.Cog): def __init__(self, bot): self.bot = bot + # Compiled recall regular expression for filtering + self._recall_filter = re.compile(self.bot.config.recall_filter, + re.IGNORECASE) + + # Default values by line number for status + self._metro_statuses = { + METRO_GREEN_LINE: METRO_NORMAL_SERVICE_MESSAGE, + METRO_ORANGE_LINE: METRO_NORMAL_SERVICE_MESSAGE, + METRO_YELLOW_LINE: METRO_NORMAL_SERVICE_MESSAGE, + METRO_BLUE_LINE: METRO_NORMAL_SERVICE_MESSAGE, + } + + self._recall_channel = None + self._metro_status_channel = None + + @commands.Cog.listener() + async def on_ready(self): + self._recall_channel = utils.get(self.bot.get_guild( + self.bot.config.server_id).text_channels, + name=self.bot.config.recall_channel) + + self._metro_status_channel = utils.get( + self.bot.get_guild(self.bot.config.server_id).text_channels, + name=self.bot.config.metro_status_channel) + + # Register all subscribers + self.bot.loop.create_task(self.cfia_rss()) + self.bot.loop.create_task(self.metro_status()) + + @canary_subscriber(12 * 3600) # run every 12 hours async def cfia_rss(self): # Written by @jidicula """ Co-routine that periodically checks the CFIA Health Hazard Alerts RSS feed for updates. """ - await self.bot.wait_until_ready() - while not self.bot.is_closed(): - recall_channel = utils.get(self.bot.get_guild( - self.bot.config.server_id).text_channels, - name=self.bot.config.recall_channel) - newest_recalls = feedparser.parse(CFIA_FEED_URL)['entries'] - try: - id_unpickle = open("pickles/recall_tag.obj", 'rb') + + newest_recalls = feedparser.parse(CFIA_FEED_URL)["entries"] + + try: + with open(CFIA_RECALL_TAG_PATH, "rb") as id_unpickle: recalls = pickle.load(id_unpickle) - id_unpickle.close() - except Exception: - recalls = {} - new_recalls = False - for recall in newest_recalls: - recall_id = recall['id'] - if recall_id not in recalls: - new_recalls = True - recalls[recall_id] = "" - recall_warning = discord.Embed(title=recall['title'], - description=recall['link']) - soup = BeautifulSoup(recall['summary'], "html.parser") - try: - img_url = soup.img['src'] - summary = soup.p.find_parent().text.strip() - except Exception: - img_url = "" - summary = recall['summary'] - if re.search(self.bot.config.recall_filter, summary, - re.IGNORECASE): - recall_warning.set_image(url=img_url) - recall_warning.add_field(name="Summary", value=summary) - await recall_channel.send(embed=recall_warning) - if new_recalls: - # Pickle newly added IDs - id_pickle = open("pickles/recall_tag.obj", 'wb') + except Exception: # TODO: Specify exception + recalls = {} + + new_recalls = False + + for recall in newest_recalls: + recall_id = recall["id"] + if recall_id in recalls: + # Don't send already-sent recalls + continue + + new_recalls = True + recalls[recall_id] = "" + recall_warning = discord.Embed(title=recall["title"], + description=recall["link"]) + soup = BeautifulSoup(recall["summary"], "html.parser") + + try: + img_url = soup.img["src"] + summary = soup.p.find_parent().text.strip() + except Exception: # TODO: Specify exception + img_url = "" + summary = recall["summary"] + + if self._recall_filter.search(summary): + recall_warning.set_image(url=img_url) + recall_warning.add_field(name="Summary", value=summary) + await self._recall_channel.send(embed=recall_warning) + + if new_recalls: + # Pickle newly added IDs + with open(CFIA_RECALL_TAG_PATH, "wb") as id_pickle: pickle.dump(recalls, id_pickle) - id_pickle.close() - await asyncio.sleep(12 * 3600) # run every 12 hours + @staticmethod + def _check_metro_status(line_number, response_data): + # Helper function to return line name and status. + # - `line_number` must be a string containing the number of the + # metro line + # - `response` must be a JSON response object from a GET request to + # the metro status API. + line_name = response_data["metro"][line_number]["name"] + status = response_data["metro"][line_number]["data"]["text"] + return line_name, status + + @canary_subscriber(60) # Run every 60 seconds async def metro_status(self): # Written by @jidicula """ Co-routine that periodically checks the STM Metro status API for outages. """ - await self.bot.wait_until_ready() - - def check_status(line_number, response): - # Helper function to return line name and status. - # - `line_number` must be a string containing the number of the - # metro line - # - `response` must be a JSON response object from a GET request to - # the metro status API. - line_name = response.json()["metro"][line_number]["name"] - status = response.json()["metro"][line_number]["data"]["text"] - return (line_name, status) - - while not self.bot.is_closed(): - metro_status_channel = utils.get( - self.bot.get_guild(self.bot.config.server_id).text_channels, - name=self.bot.config.metro_status_channel) + + try: response = requests.get(METRO_STATUS_API) - for line_status in metro_status.items(): - line_number = line_status[0] - cached_status = line_status[1] - line_name, current_status = check_status(line_number, response) - if (current_status != cached_status - and current_status != METRO_INTERIM_STATUS): - metro_status[line_number] = current_status - metro_status_update = discord.Embed( - title=line_name, - description=current_status, - colour=METRO_COLOURS[line_number]) - await metro_status_channel.send(embed=metro_status_update) - await asyncio.sleep(60) # Run every 60 seconds + response_data = response.json() + except json.decoder.JSONDecodeError: + # STM API sometimes returns non-JSON responses + return + + for line_number, cached_status in self._metro_statuses.items(): + line_name, current_status = Subscribers._check_metro_status( + line_number, response_data) + if current_status in (cached_status, METRO_INTERIM_STATUS): + # Don't send message if the status hasn't changed or the status + # is currently in the middle of changing on the API side. + continue + + self._metro_statuses[line_number] = current_status + metro_status_update = discord.Embed( + title=line_name, + description=current_status, + colour=METRO_COLOURS[line_number]) + + await self._metro_status_channel.send(embed=metro_status_update) def setup(bot): bot.add_cog(Subscribers(bot)) - bot.loop.create_task(Subscribers(bot).cfia_rss()) - bot.loop.create_task(Subscribers(bot).metro_status()) diff --git a/cogs/utils/subscribers.py b/cogs/utils/subscribers.py new file mode 100644 index 00000000..b461e952 --- /dev/null +++ b/cogs/utils/subscribers.py @@ -0,0 +1,58 @@ +# Copyright (C) idoneam (2016-2020) +# +# This file is part of Canary +# +# Canary is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Canary is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with Canary. If not, see . + +import asyncio +import functools +from typing import Callable + +__all__ = [ + "CanarySubscriberException", + "canary_subscriber", +] + + +class CanarySubscriberException(Exception): + pass + + +NO_BOT = CanarySubscriberException("Could not get bot from wrapped function") + + +def canary_subscriber(sleep_time: int): + def _canary_subscriber(func: Callable): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + if not args: + raise NO_BOT + + try: + bot = getattr(args[0], "bot") + except AttributeError: + raise NO_BOT + + await bot.wait_until_ready() + while not bot.is_closed(): + try: + await func(*args, **kwargs) + except Exception as e: + bot.logger.error("Subscriber encountered error:") + bot.log_traceback(e) + await asyncio.sleep(sleep_time) + + return wrapper + + return _canary_subscriber