diff --git a/wolfram/wolfram.py b/wolfram/wolfram.py index 1c4d298..900d82f 100644 --- a/wolfram/wolfram.py +++ b/wolfram/wolfram.py @@ -1,27 +1,30 @@ import os import aiohttp -from .utils import checks from discord.ext import commands import xml.etree.ElementTree as ET from cogs.utils.dataIO import dataIO +from .utils import checks +from .utils.chat_formatting import escape_mass_mentions +from .utils.chat_formatting import box +from __main__ import send_cmd_help class Wolfram: def __init__(self, bot): self.bot = bot - self.settings = dataIO.load_json('data/wolfram/settings.json') + self.settings = dataIO.load_json("data/wolfram/settings.json") - @commands.command(pass_context=True, name='wolfram', aliases=['ask']) - async def _wolfram(self, context, *arguments: str): + @commands.command(pass_context=True, name="wolfram", aliases=["ask"]) + async def _wolfram(self, ctx, *arguments: str): """ Ask Wolfram Alpha any question """ - api_key = self.settings['WOLFRAM_API_KEY'] + api_key = self.settings["WOLFRAM_API_KEY"] if api_key: - url = 'http://api.wolframalpha.com/v2/query?' - query = ' '.join(arguments) - payload = {'input': query, 'appid': api_key} - headers = {'user-agent': 'Red-cog/1.0.0'} + url = "http://api.wolframalpha.com/v2/query?" + query = " ".join(arguments) + payload = {"input": query, "appid": api_key} + headers = {"user-agent": "Red-cog/1.0.0"} conn = aiohttp.TCPConnector(verify_ssl=False) session = aiohttp.ClientSession(connector=conn) async with session.get(url, params=payload, headers=headers) as r: @@ -29,40 +32,46 @@ class Wolfram: session.close() root = ET.fromstring(result) a = [] - for pt in root.findall('.//plaintext'): + for pt in root.findall(".//plaintext"): if pt.text: a.append(pt.text.capitalize()) if len(a) < 1: - message = 'There is as yet insufficient data for a meaningful answer.' + message = "There is as yet insufficient data for a meaningful answer." else: - message = '\n'.join(a[0:3]) + message = "\n".join(a[0:3]) else: - message = 'No API key set for Wolfram Alpha. Get one at http://products.wolframalpha.com/api/' - await self.bot.say('```{0}```'.format(message)) + message = ( + "No API key set for Wolfram Alpha. Get one at http://products.wolframalpha.com/api/" + ) + message = escape_mass_mentions(message) + await self.bot.say(box(message)) - @commands.command(pass_context=True, name='setwolframapi', aliases=['setwolfram']) + @commands.command(pass_context=True, name="setwolframapi", aliases=["setwolfram"]) @checks.is_owner() - async def _setwolframapi(self, context, key: str): + async def _setwolframapi(self, ctx, key: str): """ Set the api-key """ if key: - self.settings['WOLFRAM_API_KEY'] = key - dataIO.save_json('data/wolfram/settings.json', self.settings) + self.settings["WOLFRAM_API_KEY"] = key + dataIO.save_json("data/wolfram/settings.json", self.settings) + await self.bot.say("Key set.") + else: + await send_cmd_help(ctx) def check_folder(): - if not os.path.exists('data/wolfram'): - print('Creating data/wolfram folder...') - os.makedirs('data/wolfram') + if not os.path.exists("data/wolfram"): + print("Creating data/wolfram folder...") + os.makedirs("data/wolfram") def check_file(): data = {} - data['WOLFRAM_API_KEY'] = False - f = 'data/wolfram/settings.json' + data["WOLFRAM_API_KEY"] = False + f = "data/wolfram/settings.json" if not dataIO.is_valid_json(f): - print('Creating default settings.json...') + print("Creating default settings.json...") dataIO.save_json(f, data)