01aa67fa9a
Escaped mass mentions in the returned message. Added slightly more friendly API key interaction and Black formatting.
83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
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)
|