Upload wolfram.py

Paddo's old unsupported Wolfram cog, now with a shiny new home
2018-06-21 09:48:35 -07:00
parent 74901f5e37
commit cbd9901857
1 changed files with 143 additions and 0 deletions
--- a/wolfram/wolfram.py
+++ b/wolfram/wolfram.py
@@ -0,0 +1,143 @@
+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 PIL import Image
+
+
+class Wolfram:
+    def __init__(self, bot):
+        self.bot = bot
+        self.settings = dataIO.load_json('data/wolfram/settings.json')
+        self.session = aiohttp.ClientSession()
+
+    def __unload(self):
+        self.session.close()
+
+    @commands.command(pass_context=True, name='wolfram', aliases=['ask'])
+    async def _wolfram(self, context, *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/'
+        await self.bot.say('```{0}```'.format(message))
+
+    @commands.command(pass_context=True, name='wolframs', aliases=['asks'])
+    async def _wolframsimple(self, ctx, *arguments : str):
+        """
+            Ask Wolfram Alpha any question (using SIMPLE API)
+        """
+        user = ctx.message.author
+        channel = ctx.message.channel
+        api_key = self.settings['WOLFRAM_API_KEY']
+        width = 800
+        max_height = 2000
+        font_size = 30
+        layout = 'labelbar'
+        background = '193555'
+        foreground = 'white'
+        units = 'metric'
+
+        if api_key:
+            query = '+'.join(arguments)
+            url = 'http://api.wolframalpha.com/v1/simple?appid={}&i={}%3F&width={}&fontsize={}&layout={}&background={}&foreground={}&units={}'.format(
+                api_key, query, width, font_size, layout, background, foreground, units)
+
+            #try:
+            filename = 'data/wolfram/{}.png'.format(user.id)
+            async with self.session.get(url) as r:
+                image = await r.content.read()
+            with open(filename,'wb') as f:
+                f.write(image)
+
+            # crop image
+            image = Image.open(filename)
+            width = image.size[0]
+            height = image.size[1]
+
+            # if too big
+            if height > max_height:
+                offset = 100
+                size_det_img = image.crop((width-offset, 0, width - offset + 1, height))
+                # print('DIMENSIONS: ', size_det_img.size)
+                size_det_img = size_det_img.convert('RGB')
+                current_color = size_det_img.getpixel((0, 0))
+                change_height = 0
+                for i in range(height):
+                    new_pixel_color = size_det_img.getpixel((0, i))
+                    # print(current_color, new_pixel_color)
+                    if current_color != new_pixel_color:
+                        if i > max_height:
+                            break
+                        change_height = i
+
+                # print('CHANGE HEIGHT: ', change_height)
+
+                img2 = image.crop((0, 0, width, change_height))
+                image = img2
+
+            image.save(filename)
+
+            await self.bot.send_file(channel, content="{}".format(user.mention), fp=filename)
+            os.remove(filename)
+            #except:
+                #await self.bot.say('Error')
+                #return
+        else:
+            await self.bot.say('No API key set for Wolfram Alpha. Get one at http://products.wolframalpha.com/api/')
+            return
+
+    @commands.command(pass_context=True, name='setwolframapi', aliases=['setwolfram'])
+    @checks.is_owner()
+    async def _setwolframapi(self, context, key: str):
+        """
+        Set the api-key
+        """
+        if key:
+            self.settings['WOLFRAM_API_KEY'] = key
+            dataIO.save_json('data/wolfram/settings.json', self.settings)
+
+
+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)