From be99660463db43d4f0d6be2f5715cb5481cf85d0 Mon Sep 17 00:00:00 2001 From: Nick80835 Date: Sat, 9 May 2020 13:15:08 -0400 Subject: [PATCH] implement a unified aiohttp client --- ubot/loader.py | 2 ++ ubot/micro_bot.py | 7 ++++++- ubot/modules/4chan.py | 12 ++---------- ubot/modules/catdog.py | 22 ++++------------------ ubot/modules/corona.py | 16 ++++++++++------ ubot/modules/danbooru.py | 8 +------- ubot/modules/evaluation.py | 9 +++++++-- ubot/modules/gelbooru.py | 8 +------- ubot/modules/nekos.py | 8 +------- ubot/modules/sankaku.py | 8 +------- ubot/modules/urbandict.py | 9 +-------- 11 files changed, 36 insertions(+), 73 deletions(-) diff --git a/ubot/loader.py b/ubot/loader.py index e25b15f..853d6bf 100644 --- a/ubot/loader.py +++ b/ubot/loader.py @@ -5,6 +5,7 @@ from importlib import import_module, reload from os.path import basename, dirname, isfile from re import escape +from aiohttp import ClientSession from telethon import events from telethon.tl.types import DocumentAttributeFilename @@ -20,6 +21,7 @@ class Loader(): self.settings = settings self.command_handler = CommandHandler(client, logger, settings) self.help_dict = {} + self.aioclient = ClientSession() self.botversion = "0.1.3" def load_all_modules(self): diff --git a/ubot/micro_bot.py b/ubot/micro_bot.py index 31206ff..46748e3 100644 --- a/ubot/micro_bot.py +++ b/ubot/micro_bot.py @@ -74,8 +74,13 @@ class MicroBot(): else: self.logger.info("Stopping client.") + await self.loader.aioclient.close() await self.client.disconnect() micro_bot = MicroBot() -micro_bot.start_microbot() + +try: + micro_bot.start_microbot() +except KeyboardInterrupt: + micro_bot.client.loop.run_until_complete(micro_bot.stop_client()) diff --git a/ubot/modules/4chan.py b/ubot/modules/4chan.py index 5beb0aa..926c213 100644 --- a/ubot/modules/4chan.py +++ b/ubot/modules/4chan.py @@ -2,8 +2,6 @@ from random import choice -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -18,29 +16,23 @@ VALID_ENDS = (".mp4", ".jpg", ".jpeg", ".png", ".gif") async def fourchan(event): as_file = bool(event.pattern_match.group(1)) - session = ClientSession() - - async with session.get(BOARD_URL.format(event.args)) as response: + async with ldr.aioclient.get(BOARD_URL.format(event.args)) as response: if response.status == 200: board_response = await response.json() op_id = choice(choice(board_response)["threads"])["no"] else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return - async with session.get(POST_URL.format(event.args, op_id)) as response: + async with ldr.aioclient.get(POST_URL.format(event.args, op_id)) as response: if response.status == 200: post_response = await response.json() post_info = choice([[i["tim"], i["ext"], i["com"] if "com" in i else None] for i in post_response["posts"] if "tim" in i and i["ext"] in VALID_ENDS]) post_file_url = CONTENT_URL.format(event.args, post_info[0], post_info[1]) else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return - await session.close() - if not response: await event.reply(f"`No results for board: `**{event.args}**") return diff --git a/ubot/modules/catdog.py b/ubot/modules/catdog.py index a8fbe14..80af7c5 100644 --- a/ubot/modules/catdog.py +++ b/ubot/modules/catdog.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: GPL-2.0-or-later -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -19,54 +17,42 @@ GIFPARAM = {"mime_types": "gif"} async def neko_atsume(params): - session = ClientSession() - - async with session.get(CAT_URL, params=params, headers=CAT_HEADERS) as response: + async with ldr.aioclient.get(CAT_URL, params=params, headers=CAT_HEADERS) as response: if response.status == 200: neko = await response.json() else: neko = response.status - await session.close() return neko async def inu_atsume(params): - session = ClientSession() - - async with session.get(DOG_URL, params=params, headers=DOG_HEADERS) as response: + async with ldr.aioclient.get(DOG_URL, params=params, headers=DOG_HEADERS) as response: if response.status == 200: inu = await response.json() else: inu = response.status - await session.close() return inu async def shibe_inu_atsume(): - session = ClientSession() - - async with session.get(SHIBE_URL, params=None, headers=None) as response: + async with ldr.aioclient.get(SHIBE_URL, params=None, headers=None) as response: if response.status == 200: shibe_inu = await response.json() else: shibe_inu = response.status - await session.close() return shibe_inu async def tori_atsume(): - session = ClientSession() - - async with session.get(BIRD_URL, params=None, headers=None) as response: + async with ldr.aioclient.get(BIRD_URL, params=None, headers=None) as response: if response.status == 200: tori = await response.json() else: tori = response.status - await session.close() return tori diff --git a/ubot/modules/corona.py b/ubot/modules/corona.py index 2573695..b4cf712 100644 --- a/ubot/modules/corona.py +++ b/ubot/modules/corona.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: GPL-2.0-or-later -from requests import get - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -10,9 +8,9 @@ ldr = micro_bot.loader @ldr.add("corona") async def corona(event): if event.args: - with get(f"https://corona.lmao.ninja/v2/countries/{event.args}") as response: - if response.status_code == 200: - response = response.json() + async with ldr.aioclient.get(f"https://corona.lmao.ninja/v2/countries/{event.args}") as response: + if response.status == 200: + response = await response.json() else: await event.reply(f"`An error occurred, response code: `**{response.status_code}**") return @@ -20,6 +18,12 @@ async def corona(event): response_text = f"`Stats for `**{response['country']}**\n\n`Cases: `**{response['cases']}** **({response['todayCases']} today)**\n`Deaths: `**{response['deaths']}** **({response['todayDeaths']} today)**\n`Recoveries: `**{response['recovered']}**" await event.reply(response_text) else: - response = get("https://corona.lmao.ninja/v2/all").json() + async with ldr.aioclient.get(f"https://corona.lmao.ninja/v2/all") as response: + if response.status == 200: + response = await response.json() + else: + await event.reply(f"`An error occurred, response code: `**{response.status_code}**") + return + response_text = f"`Global stats`\n\n`Cases: `**{response['cases']}** **({response['todayCases']} today)**\n`Deaths: `**{response['deaths']}** **({response['todayDeaths']} today)**\n`Recoveries: `**{response['recovered']}**" await event.reply(response_text) diff --git a/ubot/modules/danbooru.py b/ubot/modules/danbooru.py index dcd3a18..929d2d0 100644 --- a/ubot/modules/danbooru.py +++ b/ubot/modules/danbooru.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: GPL-2.0-or-later -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -26,15 +24,11 @@ async def danbooru(event): "random": "true", "tags": f"{rating} {event.args}".strip().replace(" ", " ")} - session = ClientSession() - - async with session.get(DAN_URL, params=params) as response: + async with ldr.aioclient.get(DAN_URL, params=params) as response: if response.status == 200: response = await response.json() - await session.close() else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return if not response: diff --git a/ubot/modules/evaluation.py b/ubot/modules/evaluation.py index fd81a75..b70aa19 100644 --- a/ubot/modules/evaluation.py +++ b/ubot/modules/evaluation.py @@ -6,7 +6,6 @@ from re import sub from gtts import gTTS from PIL import Image -from requests import get from speedtest import Speedtest from ubot.micro_bot import micro_bot @@ -79,7 +78,13 @@ async def ip_lookup(event): await event.reply("`Provide an IP!`") return - lookup_json = get(f"http://ip-api.com/json/{ip}").json() + async with ldr.aioclient.get(f"http://ip-api.com/json/{ip}") as response: + if response.status == 200: + lookup_json = await response.json() + else: + await event.reply(f"`An error occurred when looking for `**{ip}**`: `**{response.status}**") + return + fixed_lookup = {} for key, value in lookup_json.items(): diff --git a/ubot/modules/gelbooru.py b/ubot/modules/gelbooru.py index 6f8aefc..c910694 100644 --- a/ubot/modules/gelbooru.py +++ b/ubot/modules/gelbooru.py @@ -2,8 +2,6 @@ from random import choice -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -30,15 +28,11 @@ async def gelbooru(event): "json": 1, "tags": f"{rating} {event.args} sort:random".strip().replace(" ", " ")} - session = ClientSession() - - async with session.get(GEL_URL, params=params) as response: + async with ldr.aioclient.get(GEL_URL, params=params) as response: if response.status == 200: response = await response.json() - await session.close() else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return if not response: diff --git a/ubot/modules/nekos.py b/ubot/modules/nekos.py index ae3ca99..41fd91a 100644 --- a/ubot/modules/nekos.py +++ b/ubot/modules/nekos.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: GPL-2.0-or-later -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -21,15 +19,11 @@ async def supernekoatsume(event): else: reply_to = event - session = ClientSession() - - async with session.get(NEKO_URL + nekotype) as response: + async with ldr.aioclient.get(NEKO_URL + nekotype) as response: if response.status == 200: image_url = (await response.json())["url"] - await session.close() else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return try: diff --git a/ubot/modules/sankaku.py b/ubot/modules/sankaku.py index 53e2529..69868fa 100644 --- a/ubot/modules/sankaku.py +++ b/ubot/modules/sankaku.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: GPL-2.0-or-later -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -26,15 +24,11 @@ async def sankaku(event): "limit": 5, "tags": f"order:random {rating} {event.args}".strip().replace(" ", " ")} - session = ClientSession() - - async with session.get(SAN_URL, params=params) as response: + async with ldr.aioclient.get(SAN_URL, params=params) as response: if response.status == 200: response = await response.json() - await session.close() else: await event.reply(f"`An error occurred, response code: `**{response.status}**") - await session.close() return if not response: diff --git a/ubot/modules/urbandict.py b/ubot/modules/urbandict.py index 8e8ac1d..3f3efcc 100644 --- a/ubot/modules/urbandict.py +++ b/ubot/modules/urbandict.py @@ -2,8 +2,6 @@ import io -from aiohttp import ClientSession - from ubot.micro_bot import micro_bot ldr = micro_bot.loader @@ -23,18 +21,13 @@ async def urban_dict(event): params = None url = UD_RANDOM_URL - session = ClientSession() - - async with session.get(url, params=params) as response: + async with ldr.aioclient.get(url, params=params) as response: if response.status == 200: response = await response.json() else: await event.reply(f"`An error occurred, response code:` **{response.status}**") - await session.close() return - await session.close() - if response['list']: response_word = response['list'][0] else: