diff options
Diffstat (limited to '')
| -rw-r--r-- | bot/exts/evergreen/wolfram.py | 82 | 
1 files changed, 44 insertions, 38 deletions
| diff --git a/bot/exts/evergreen/wolfram.py b/bot/exts/evergreen/wolfram.py index 437d9e1a..26674d37 100644 --- a/bot/exts/evergreen/wolfram.py +++ b/bot/exts/evergreen/wolfram.py @@ -1,7 +1,7 @@  import logging  from io import BytesIO  from typing import Callable, List, Optional, Tuple -from urllib import parse +from urllib.parse import urlencode  import arrow  import discord @@ -9,6 +9,7 @@ from discord import Embed  from discord.ext import commands  from discord.ext.commands import BucketType, Cog, Context, check, group +from bot.bot import Bot  from bot.constants import Colours, STAFF_ROLES, Wolfram  from bot.utils.pagination import ImagePaginator @@ -16,7 +17,7 @@ log = logging.getLogger(__name__)  APPID = Wolfram.key  DEFAULT_OUTPUT_FORMAT = "JSON" -QUERY = "http://api.wolframalpha.com/v2/{request}?{data}" +QUERY = "http://api.wolframalpha.com/v2/{request}"  WOLF_IMAGE = "https://www.symbols.com/gi.php?type=1&id=2886&i=1"  MAX_PODS = 20 @@ -39,9 +40,11 @@ async def send_embed(      """Generate & send a response embed with Wolfram as the author."""      embed = Embed(colour=colour)      embed.description = message_txt -    embed.set_author(name="Wolfram Alpha", -                     icon_url=WOLF_IMAGE, -                     url="https://www.wolframalpha.com/") +    embed.set_author( +        name="Wolfram Alpha", +        icon_url=WOLF_IMAGE, +        url="https://www.wolframalpha.com/" +    )      if footer:          embed.set_footer(text=footer) @@ -55,14 +58,15 @@ def custom_cooldown(*ignore: List[int]) -> Callable:      """      Implement per-user and per-guild cooldowns for requests to the Wolfram API. -    A list of roles may be provided to ignore the per-user cooldown +    A list of roles may be provided to ignore the per-user cooldown.      """      async def predicate(ctx: Context) -> bool: -        if ctx.invoked_with == 'help': +        if ctx.invoked_with == "help":              # if the invoked command is help we don't want to increase the ratelimits since it's not actually              # invoking the command/making a request, so instead just check if the user/guild are on cooldown.              guild_cooldown = not guildcd.get_bucket(ctx.message).get_tokens() == 0  # if guild is on cooldown -            if not any(r.id in ignore for r in ctx.author.roles):  # check user bucket if user is not ignored +            # check the message is in a guild, and check user bucket if user is not ignored +            if ctx.guild and not any(r.id in ignore for r in ctx.author.roles):                  return guild_cooldown and not usercd.get_bucket(ctx.message).get_tokens() == 0              return guild_cooldown @@ -101,10 +105,10 @@ def custom_cooldown(*ignore: List[int]) -> Callable:      return check(predicate) -async def get_pod_pages(ctx: Context, bot: commands.Bot, query: str) -> Optional[List[Tuple]]: +async def get_pod_pages(ctx: Context, bot: Bot, query: str) -> Optional[List[Tuple]]:      """Get the Wolfram API pod pages for the provided query.""" -    async with ctx.channel.typing(): -        url_str = parse.urlencode({ +    async with ctx.typing(): +        params = {              "input": query,              "appid": APPID,              "output": DEFAULT_OUTPUT_FORMAT, @@ -112,27 +116,27 @@ async def get_pod_pages(ctx: Context, bot: commands.Bot, query: str) -> Optional              "location": "the moon",              "latlong": "0.0,0.0",              "ip": "1.1.1.1" -        }) -        request_url = QUERY.format(request="query", data=url_str) +        } +        request_url = QUERY.format(request="query") -        async with bot.http_session.get(request_url) as response: -            json = await response.json(content_type='text/plain') +        async with bot.http_session.get(url=request_url, params=params) as response: +            json = await response.json(content_type="text/plain")          result = json["queryresult"] - +        log_full_url = f"{request_url}?{urlencode(params)}"          if result["error"]:              # API key not set up correctly              if result["error"]["msg"] == "Invalid appid":                  message = "Wolfram API key is invalid or missing."                  log.warning(                      "API key seems to be missing, or invalid when " -                    f"processing a wolfram request: {url_str}, Response: {json}" +                    f"processing a wolfram request: {log_full_url}, Response: {json}"                  )                  await send_embed(ctx, message)                  return              message = "Something went wrong internally with your request, please notify staff!" -            log.warning(f"Something went wrong getting a response from wolfram: {url_str}, Response: {json}") +            log.warning(f"Something went wrong getting a response from wolfram: {log_full_url}, Response: {json}")              await send_embed(ctx, message)              return @@ -161,25 +165,25 @@ async def get_pod_pages(ctx: Context, bot: commands.Bot, query: str) -> Optional  class Wolfram(Cog):      """Commands for interacting with the Wolfram|Alpha API.""" -    def __init__(self, bot: commands.Bot): +    def __init__(self, bot: Bot):          self.bot = bot      @group(name="wolfram", aliases=("wolf", "wa"), invoke_without_command=True)      @custom_cooldown(*STAFF_ROLES)      async def wolfram_command(self, ctx: Context, *, query: str) -> None:          """Requests all answers on a single image, sends an image of all related pods.""" -        url_str = parse.urlencode({ +        params = {              "i": query,              "appid": APPID,              "location": "the moon",              "latlong": "0.0,0.0",              "ip": "1.1.1.1" -        }) -        query = QUERY.format(request="simple", data=url_str) +        } +        request_url = QUERY.format(request="simple")          # Give feedback that the bot is working. -        async with ctx.channel.typing(): -            async with self.bot.http_session.get(query) as response: +        async with ctx.typing(): +            async with self.bot.http_session.get(url=request_url, params=params) as response:                  status = response.status                  image_bytes = await response.read() @@ -187,11 +191,11 @@ class Wolfram(Cog):              image_url = "attachment://image.png"              if status == 501: -                message = "Failed to get response" +                message = "Failed to get response."                  footer = ""                  color = Colours.soft_red              elif status == 400: -                message = "No input found" +                message = "No input found."                  footer = ""                  color = Colours.soft_red              elif status == 403: @@ -220,9 +224,11 @@ class Wolfram(Cog):              return          embed = Embed() -        embed.set_author(name="Wolfram Alpha", -                         icon_url=WOLF_IMAGE, -                         url="https://www.wolframalpha.com/") +        embed.set_author( +            name="Wolfram Alpha", +            icon_url=WOLF_IMAGE, +            url="https://www.wolframalpha.com/" +        )          embed.colour = Colours.soft_orange          await ImagePaginator.paginate(pages, ctx, embed) @@ -251,28 +257,28 @@ class Wolfram(Cog):      @custom_cooldown(*STAFF_ROLES)      async def wolfram_short_command(self, ctx: Context, *, query: str) -> None:          """Requests an answer to a simple question.""" -        url_str = parse.urlencode({ +        params = {              "i": query,              "appid": APPID,              "location": "the moon",              "latlong": "0.0,0.0",              "ip": "1.1.1.1" -        }) -        query = QUERY.format(request="result", data=url_str) +        } +        request_url = QUERY.format(request="result")          # Give feedback that the bot is working. -        async with ctx.channel.typing(): -            async with self.bot.http_session.get(query) as response: +        async with ctx.typing(): +            async with self.bot.http_session.get(url=request_url, params=params) as response:                  status = response.status                  response_text = await response.text()              if status == 501: -                message = "Failed to get response" +                message = "Failed to get response."                  color = Colours.soft_red              elif status == 400: -                message = "No input found" +                message = "No input found."                  color = Colours.soft_red -            elif response_text == "Error 1: Invalid appid": +            elif response_text == "Error 1: Invalid appid.":                  message = "Wolfram API key is invalid or missing."                  color = Colours.soft_red              else: @@ -282,6 +288,6 @@ class Wolfram(Cog):              await send_embed(ctx, message, color) -def setup(bot: commands.Bot) -> None: +def setup(bot: Bot) -> None:      """Load the Wolfram cog."""      bot.add_cog(Wolfram(bot)) | 
