Merge pull request #131 from classabbyamp/aiohttp-sessions

combine all aiohttp operations into a single session

common.py

@@ -24,7 +24,7 @@ import discord.ext.commands as commands
 import data.options as opt
 
 
-__all__ = ["colours", "cat", "emojis", "error_embed_factory", "add_react", "check_if_owner"]
+__all__ = ["colours", "cat", "emojis", "embed_factory", "error_embed_factory", "add_react", "check_if_owner"]
 
 
 # --- Common values ---

exts/ae7q.py

@@ -16,7 +16,6 @@ NA2AAA: unassigned, no records
 import discord.ext.commands as commands
 
 from bs4 import BeautifulSoup
-import aiohttp
 
 import common as cmn
 
@@ -24,6 +23,7 @@ import common as cmn
 class AE7QCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
+        self.session = bot.qrm.session
 
     @commands.group(name="ae7q", aliases=["ae"], category=cmn.cat.lookup)
     async def _ae7q_lookup(self, ctx: commands.Context):
@@ -39,15 +39,14 @@ class AE7QCog(commands.Cog):
         base_url = "http://ae7q.com/query/data/CallHistory.php?CALL="
         embed = cmn.embed_factory(ctx)
 
-        async with aiohttp.ClientSession() as session:
+        async with self.session.get(base_url + callsign) as resp:
-            async with session.get(base_url + callsign) as resp:
+            if resp.status != 200:
-                if resp.status != 200:
+                embed.title = "Error in AE7Q call command"
-                    embed.title = "Error in AE7Q call command"
+                embed.description = 'Could not load AE7Q'
-                    embed.description = 'Could not load AE7Q'
+                embed.colour = cmn.colours.bad
-                    embed.colour = cmn.colours.bad
+                await ctx.send(embed=embed)
-                    await ctx.send(embed=embed)
+                return
-                    return
+            page = await resp.text()
-                page = await resp.text()
 
         soup = BeautifulSoup(page, features="html.parser")
         tables = soup.select("table.Database")

exts/image.py

@@ -12,14 +12,13 @@ import io
 import discord
 import discord.ext.commands as commands
 
-import aiohttp
-
 import common as cmn
 
 
 class ImageCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
+        self.session = bot.qrm.session
 
     @commands.command(name="bandplan", aliases=['plan', 'bands'], category=cmn.cat.ref)
     async def _bandplan(self, ctx: commands.Context, region: str = ''):
@@ -58,14 +57,13 @@ class ImageCog(commands.Cog):
             embed = cmn.embed_factory(ctx)
             embed.title = 'Current Greyline Conditions'
             embed.colour = cmn.colours.good
-            async with aiohttp.ClientSession() as session:
+            async with self.session.get(gl_url) as resp:
-                async with session.get(gl_url) as resp:
+                if resp.status != 200:
-                    if resp.status != 200:
+                    embed.description = 'Could not download file...'
-                        embed.description = 'Could not download file...'
+                    embed.colour = cmn.colours.bad
-                        embed.colour = cmn.colours.bad
+                else:
-                    else:
+                    data = io.BytesIO(await resp.read())
-                        data = io.BytesIO(await resp.read())
+                    embed.set_image(url=f'attachment://greyline.jpg')
-                        embed.set_image(url=f'attachment://greyline.jpg')
         await ctx.send(embed=embed, file=discord.File(data, 'greyline.jpg'))
 
     @commands.command(name="map", category=cmn.cat.maps)

exts/qrz.py

@@ -21,7 +21,7 @@ import data.keys as keys
 class QRZCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
-        self.session = aiohttp.ClientSession()
+        self.session = bot.qrm.session
         self._qrz_session_init.start()
 
     @commands.command(name="call", aliases=["qrz"], category=cmn.cat.lookup)

exts/study.py

@@ -12,8 +12,6 @@ import json
 
 import discord.ext.commands as commands
 
-import aiohttp
-
 import common as cmn
 
 
@@ -22,6 +20,7 @@ class StudyCog(commands.Cog):
         self.bot = bot
         self.lastq = dict()
         self.source = 'Data courtesy of [HamStudy.org](https://hamstudy.org/)'
+        self.session = bot.qrm.session
 
     @commands.command(name="hamstudy", aliases=['rq', 'randomquestion', 'randomq'], category=cmn.cat.study)
     async def _random_question(self, ctx: commands.Context, level: str = None):
@@ -58,15 +57,14 @@ class StudyCog(commands.Cog):
                 await ctx.send(embed=embed)
                 return
 
-            async with aiohttp.ClientSession() as session:
+            async with self.session.get(f'https://hamstudy.org/pools/{selected_pool}') as resp:
-                async with session.get(f'https://hamstudy.org/pools/{selected_pool}') as resp:
+                if resp.status != 200:
-                    if resp.status != 200:
+                    embed.title = 'Error in HamStudy command'
-                        embed.title = 'Error in HamStudy command'
+                    embed.description = 'Could not load questions'
-                        embed.description = 'Could not load questions'
+                    embed.colour = cmn.colours.bad
-                        embed.colour = cmn.colours.bad
+                    await ctx.send(embed=embed)
-                        await ctx.send(embed=embed)
+                    return
-                        return
+                pool = json.loads(await resp.read())['pool']
-                    pool = json.loads(await resp.read())['pool']
 
             # Select a question
             pool_section = random.choice(pool)['sections']

exts/weather.py

@@ -13,8 +13,6 @@ import re
 import discord
 import discord.ext.commands as commands
 
-import aiohttp
-
 import common as cmn
 
 
@@ -23,6 +21,7 @@ class WeatherCog(commands.Cog):
 
     def __init__(self, bot: commands.Bot):
         self.bot = bot
+        self.session = bot.qrm.session
 
     @commands.command(name="bandconditions", aliases=['cond', 'condx', 'conditions'], category=cmn.cat.weather)
     async def _band_conditions(self, ctx: commands.Context):
@@ -31,14 +30,13 @@ class WeatherCog(commands.Cog):
             embed = cmn.embed_factory(ctx)
             embed.title = 'Current Solar Conditions'
             embed.colour = cmn.colours.good
-            async with aiohttp.ClientSession() as session:
+            async with self.session.get('http://www.hamqsl.com/solarsun.php') as resp:
-                async with session.get('http://www.hamqsl.com/solarsun.php') as resp:
+                if resp.status != 200:
-                    if resp.status != 200:
+                    embed.description = 'Could not download file...'
-                        embed.description = 'Could not download file...'
+                    embed.colour = cmn.colours.bad
-                        embed.colour = cmn.colours.bad
+                else:
-                    else:
+                    data = io.BytesIO(await resp.read())
-                        data = io.BytesIO(await resp.read())
+                    embed.set_image(url=f'attachment://condx.png')
-                        embed.set_image(url=f'attachment://condx.png')
         await ctx.send(embed=embed, file=discord.File(data, 'condx.png'))
 
     @commands.group(name="weather", aliases=['wttr'], category=cmn.cat.weather)
@@ -81,14 +79,13 @@ See help for weather command for possible location types. Add a `-c` or `-f` to
             embed.colour = cmn.colours.good
 
             loc = loc.replace(' ', '+')
-            async with aiohttp.ClientSession() as session:
+            async with self.session.get(f'http://wttr.in/{loc}_{units}pnFQ.png') as resp:
-                async with session.get(f'http://wttr.in/{loc}_{units}pnFQ.png') as resp:
+                if resp.status != 200:
-                    if resp.status != 200:
+                    embed.description = 'Could not download file...'
-                        embed.description = 'Could not download file...'
+                    embed.colour = cmn.colours.bad
-                        embed.colour = cmn.colours.bad
+                else:
-                    else:
+                    data = io.BytesIO(await resp.read())
-                        data = io.BytesIO(await resp.read())
+                    embed.set_image(url=f'attachment://wttr_forecast.png')
-                        embed.set_image(url=f'attachment://wttr_forecast.png')
         await ctx.send(embed=embed, file=discord.File(data, f'wttr_forecast.png'))
 
     @_weather_conditions.command(name='now', aliases=['n'], category=cmn.cat.weather)
@@ -115,14 +112,13 @@ See help for weather command for possible location types. Add a `-c` or `-f` to
             embed.colour = cmn.colours.good
 
             loc = loc.replace(' ', '+')
-            async with aiohttp.ClientSession() as session:
+            async with self.session.get(f'http://wttr.in/{loc}_0{units}pnFQ.png') as resp:
-                async with session.get(f'http://wttr.in/{loc}_0{units}pnFQ.png') as resp:
+                if resp.status != 200:
-                    if resp.status != 200:
+                    embed.description = 'Could not download file...'
-                        embed.description = 'Could not download file...'
+                    embed.colour = cmn.colours.bad
-                        embed.colour = cmn.colours.bad
+                else:
-                    else:
+                    data = io.BytesIO(await resp.read())
-                        data = io.BytesIO(await resp.read())
+                    embed.set_image(url=f'attachment://wttr_now.png')
-                        embed.set_image(url=f'attachment://wttr_now.png')
         await ctx.send(embed=embed, file=discord.File(data, 'wttr_now.png'))
 
 

main.py

@@ -10,8 +10,10 @@ General Public License, version 2.
 
 from datetime import time, datetime
 import random
+from types import SimpleNamespace
 
 import pytz
+import aiohttp
 
 import discord
 from discord.ext import commands, tasks
@@ -37,6 +39,9 @@ bot = commands.Bot(command_prefix=opt.prefix,
                    description=info.description,
                    help_command=commands.MinimalHelpCommand())
 
+bot.qrm = SimpleNamespace()
+bot.qrm.session = aiohttp.ClientSession(headers={'User-Agent': f'discord-qrm2/{info.release}'})
+
 
 # --- Commands ---
 
@@ -44,6 +49,7 @@ bot = commands.Bot(command_prefix=opt.prefix,
 @commands.check(cmn.check_if_owner)
 async def _restart_bot(ctx: commands.Context):
     """Restarts the bot."""
+    await bot.qrm.session.close()
     global exit_code
     await cmn.add_react(ctx.message, cmn.emojis.good)
     print(f"[**] Restarting! Requested by {ctx.author}.")
@@ -55,6 +61,7 @@ async def _restart_bot(ctx: commands.Context):
 @commands.check(cmn.check_if_owner)
 async def _shutdown_bot(ctx: commands.Context):
     """Shuts down the bot."""
+    await bot.qrm.session.close()
     global exit_code
     await cmn.add_react(ctx.message, cmn.emojis.good)
     print(f"[**] Shutting down! Requested by {ctx.author}.")