import os import aiohttp 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") @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"] 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"} conn = aiohttp.TCPConnector(verify_ssl=False) session = aiohttp.ClientSession(connector=conn) async with session.get(url, params=payload, headers=headers) as r: result = await r.text() session.close() root = ET.fromstring(result) a = [] 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." else: message = "\n".join(a[0:3]) else: 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"]) @checks.is_owner() 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) 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") def check_file(): data = {} data["WOLFRAM_API_KEY"] = False f = "data/wolfram/settings.json" if not dataIO.is_valid_json(f): print("Creating default settings.json...") dataIO.save_json(f, data) def setup(bot): check_folder() check_file() n = Wolfram(bot) bot.add_cog(n)