[V2 Wolfram] Escape mass mentions

Escaped mass mentions in the returned message. Added slightly more friendly API key interaction and Black formatting.
This commit is contained in:
aikaterna
2018-09-29 00:42:44 -07:00
parent adb3efa746
commit 01aa67fa9a

View File

@@ -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)