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..f98c869 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,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}.")