From abc62d3e09a39decfdd400ed6f2e9c631ef2eca1 Mon Sep 17 00:00:00 2001
From: kwzrd <kwzrd@pythondiscord.com>
Date: Sun, 29 Mar 2020 14:29:36 +0200
Subject: Deseasonify: implement generic media asset setter

Current `set_avatar`, `set_icon` and `set_banner` methods are almost
identical - they only differ in the type of asset they upload. This
leads to a lot of code repetition, especially w.r.t. error handling.

We instead add a generic method that is parametrized by an AssetType
param, and by the target entity (i.e. bot, or guild) that the asset
should be applied to. All error handling can then be done in one place.

Error handling methodology is adjusted - instead of suppressing errors,
we catch and log them. Since we no longer determine whether the upload
succeeded based on 'before' != 'after', this solves a bug where
re-applying the same asset resulted in a warning-level log, triggering
Sentry.
---
 bot/bot.py | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

(limited to 'bot')

diff --git a/bot/bot.py b/bot/bot.py
index 76ce5607..0809f7a5 100644
--- a/bot/bot.py
+++ b/bot/bot.py
@@ -2,12 +2,12 @@ import asyncio
 import enum
 import logging
 import socket
-from typing import Optional
+from typing import Optional, Union
 
 import async_timeout
 import discord
 from aiohttp import AsyncResolver, ClientSession, TCPConnector
-from discord import DiscordException, Embed
+from discord import DiscordException, Embed, Guild, User
 from discord.ext import commands
 
 from bot.constants import Channels, Client
@@ -85,6 +85,34 @@ class SeasonalBot(commands.Bot):
             return None
         return guild.me
 
+    async def _apply_asset(self, target: Union[Guild, User], asset: AssetType, url: str) -> bool:
+        """
+        Internal method for applying media assets to the guild or the bot.
+
+        This shouldn't be called directly. The purpose of this method is mainly generic
+        error handling to reduce needless code repetition.
+
+        Return True if upload was successful, False otherwise.
+        """
+        log.info(f"Attempting to set {asset.name}: {url}")
+
+        kwargs = {asset.value: await self._fetch_image(url)}
+        try:
+            async with async_timeout.timeout(5):
+                await target.edit(**kwargs)
+
+        except asyncio.TimeoutError:
+            log.info("Asset upload timed out")
+            return False
+
+        except discord.HTTPException as discord_error:
+            log.exception("Asset upload failed", exc_info=discord_error)
+            return False
+
+        else:
+            log.info(f"Asset successfully applied")
+            return True
+
     @mock_in_debug(return_value=True)
     async def set_avatar(self, url: str) -> bool:
         """Sets the bot's avatar based on a URL."""
-- 
cgit v1.2.3