From f1993c85b2ef7ecb6811b3030c73dc705fa8e76e Mon Sep 17 00:00:00 2001 From: Abigail Date: Mon, 23 Dec 2019 17:54:20 -0500 Subject: [PATCH 1/2] combine all aiohttp operations into a single session Fixes #49 --- common.py | 2 +- exts/ae7q.py | 19 +++++++++---------- exts/image.py | 18 ++++++++---------- exts/qrz.py | 2 +- exts/study.py | 20 +++++++++----------- exts/weather.py | 48 ++++++++++++++++++++++-------------------------- main.py | 7 +++++++ 7 files changed, 57 insertions(+), 59 deletions(-) diff --git a/common.py b/common.py index 494762c..b0f08fa 100644 --- a/common.py +++ b/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 --- diff --git a/exts/ae7q.py b/exts/ae7q.py index b543b13..e27a019 100644 --- a/exts/ae7q.py +++ b/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 session.get(base_url + callsign) as resp: - if resp.status != 200: - embed.title = "Error in AE7Q call command" - embed.description = 'Could not load AE7Q' - embed.colour = cmn.colours.bad - await ctx.send(embed=embed) - return - page = await resp.text() + async with self.session.get(base_url + callsign) as resp: + if resp.status != 200: + embed.title = "Error in AE7Q call command" + embed.description = 'Could not load AE7Q' + embed.colour = cmn.colours.bad + await ctx.send(embed=embed) + return + page = await resp.text() soup = BeautifulSoup(page, features="html.parser") tables = soup.select("table.Database") diff --git a/exts/image.py b/exts/image.py index f9ebd8b..5e6e048 100644 --- a/exts/image.py +++ b/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 session.get(gl_url) as resp: - if resp.status != 200: - embed.description = 'Could not download file...' - embed.colour = cmn.colours.bad - else: - data = io.BytesIO(await resp.read()) - embed.set_image(url=f'attachment://greyline.jpg') + async with self.session.get(gl_url) as resp: + if resp.status != 200: + embed.description = 'Could not download file...' + embed.colour = cmn.colours.bad + else: + data = io.BytesIO(await resp.read()) + 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) diff --git a/exts/qrz.py b/exts/qrz.py index 5846062..27831a1 100644 --- a/exts/qrz.py +++ b/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) diff --git a/exts/study.py b/exts/study.py index 143f9cd..6c98fb4 100644 --- a/exts/study.py +++ b/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 session.get(f'https://hamstudy.org/pools/{selected_pool}') as resp: - if resp.status != 200: - embed.title = 'Error in HamStudy command' - embed.description = 'Could not load questions' - embed.colour = cmn.colours.bad - await ctx.send(embed=embed) - return - pool = json.loads(await resp.read())['pool'] + async with self.session.get(f'https://hamstudy.org/pools/{selected_pool}') as resp: + if resp.status != 200: + embed.title = 'Error in HamStudy command' + embed.description = 'Could not load questions' + embed.colour = cmn.colours.bad + await ctx.send(embed=embed) + return + pool = json.loads(await resp.read())['pool'] # Select a question pool_section = random.choice(pool)['sections'] diff --git a/exts/weather.py b/exts/weather.py index 577e495..ec38443 100644 --- a/exts/weather.py +++ b/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 session.get('http://www.hamqsl.com/solarsun.php') as resp: - if resp.status != 200: - embed.description = 'Could not download file...' - embed.colour = cmn.colours.bad - else: - data = io.BytesIO(await resp.read()) - embed.set_image(url=f'attachment://condx.png') + async with self.session.get('http://www.hamqsl.com/solarsun.php') as resp: + if resp.status != 200: + embed.description = 'Could not download file...' + embed.colour = cmn.colours.bad + else: + data = io.BytesIO(await resp.read()) + 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 session.get(f'http://wttr.in/{loc}_{units}pnFQ.png') as resp: - if resp.status != 200: - embed.description = 'Could not download file...' - embed.colour = cmn.colours.bad - else: - data = io.BytesIO(await resp.read()) - embed.set_image(url=f'attachment://wttr_forecast.png') + async with self.session.get(f'http://wttr.in/{loc}_{units}pnFQ.png') as resp: + if resp.status != 200: + embed.description = 'Could not download file...' + embed.colour = cmn.colours.bad + else: + data = io.BytesIO(await resp.read()) + 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 session.get(f'http://wttr.in/{loc}_0{units}pnFQ.png') as resp: - if resp.status != 200: - embed.description = 'Could not download file...' - embed.colour = cmn.colours.bad - else: - data = io.BytesIO(await resp.read()) - embed.set_image(url=f'attachment://wttr_now.png') + async with self.session.get(f'http://wttr.in/{loc}_0{units}pnFQ.png') as resp: + if resp.status != 200: + embed.description = 'Could not download file...' + embed.colour = cmn.colours.bad + else: + data = io.BytesIO(await resp.read()) + embed.set_image(url=f'attachment://wttr_now.png') await ctx.send(embed=embed, file=discord.File(data, 'wttr_now.png')) diff --git a/main.py b/main.py index a7e7b7b..3a110a6 100644 --- a/main.py +++ b/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,13 +39,17 @@ 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 --- + @bot.command(name="restart", hidden=True) @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}.") From f7aaa467b5e87ed70f3183ff7199323bb78813da Mon Sep 17 00:00:00 2001 From: Abigail Date: Mon, 23 Dec 2019 18:10:36 -0500 Subject: [PATCH 2/2] >_> pep8 --- main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/main.py b/main.py index 3a110a6..f98c869 100644 --- a/main.py +++ b/main.py @@ -42,8 +42,8 @@ bot = commands.Bot(command_prefix=opt.prefix, bot.qrm = SimpleNamespace() bot.qrm.session = aiohttp.ClientSession(headers={'User-Agent': f'discord-qrm2/{info.release}'}) -# --- Commands --- +# --- Commands --- @bot.command(name="restart", hidden=True) @commands.check(cmn.check_if_owner)