Merge pull request #483 from WillDaSilva/prideavatar-url

Make prideavatar support specifying the image by URL

1 files changed, 70 insertions, 37 deletions

diff --git a/bot/exts/pride/pride_avatar.py b/bot/exts/pride/pride_avatar.py
index 3f9878e3..2eade796 100644
--- a/bot/exts/pride/pride_avatar.py
+++ b/bot/exts/pride/pride_avatar.py
@@ -1,10 +1,12 @@
 import logging
 from io import BytesIO
 from pathlib import Path
+from typing import Tuple
 
+import aiohttp
 import discord
-from PIL import Image, ImageDraw
-from discord.ext import commands
+from PIL import Image, ImageDraw, UnidentifiedImageError
+from discord.ext.commands import Bot, Cog, Context, group
 
 from bot.constants import Colours
 
@@ -53,10 +55,10 @@ OPTIONS = {
 }
 
 
-class PrideAvatar(commands.Cog):
+class PrideAvatar(Cog):
     """Put an LGBT spin on your avatar!"""
 
-    def __init__(self, bot: commands.Bot):
+    def __init__(self, bot: Bot):
         self.bot = bot
 
     @staticmethod
@@ -78,8 +80,41 @@ class PrideAvatar(commands.Cog):
         ring.putalpha(mask)
         return ring
 
-    @commands.group(aliases=["avatarpride", "pridepfp", "prideprofile"], invoke_without_command=True)
-    async def prideavatar(self, ctx: commands.Context, option: str = "lgbt", pixels: int = 64) -> None:
+    @staticmethod
+    def process_options(option: str, pixels: int) -> Tuple[str, int, str]:
+        """Does some shared preprocessing for the prideavatar commands."""
+        return option.lower(), max(0, min(512, pixels)), OPTIONS.get(option)
+
+    async def process_image(self, ctx: Context, image_bytes: bytes, pixels: int, flag: str, option: str) -> None:
+        """Constructs the final image, embeds it, and sends it."""
+        try:
+            avatar = Image.open(BytesIO(image_bytes))
+        except UnidentifiedImageError:
+            return await ctx.send("Cannot identify image from provided URL")
+        avatar = avatar.convert("RGBA").resize((1024, 1024))
+
+        avatar = self.crop_avatar(avatar)
+
+        ring = Image.open(Path(f"bot/resources/pride/flags/{flag}.png")).resize((1024, 1024))
+        ring = ring.convert("RGBA")
+        ring = self.crop_ring(ring, pixels)
+
+        avatar.alpha_composite(ring, (0, 0))
+        bufferedio = BytesIO()
+        avatar.save(bufferedio, format="PNG")
+        bufferedio.seek(0)
+
+        file = discord.File(bufferedio, filename="pride_avatar.png")  # Creates file to be used in embed
+        embed = discord.Embed(
+            name="Your Lovely Pride Avatar",
+            description=f"Here is your lovely avatar, surrounded by\n a beautiful {option} flag. Enjoy :D"
+        )
+        embed.set_image(url="attachment://pride_avatar.png")
+        embed.set_footer(text=f"Made by {ctx.author.display_name}", icon_url=ctx.author.avatar_url)
+        await ctx.send(file=file, embed=embed)
+
+    @group(aliases=["avatarpride", "pridepfp", "prideprofile"], invoke_without_command=True)
+    async def prideavatar(self, ctx: Context, option: str = "lgbt", pixels: int = 64) -> None:
         """
         This surrounds an avatar with a border of a specified LGBT flag.
 
@@ -88,45 +123,43 @@ class PrideAvatar(commands.Cog):
         This has a maximum of 512px and defaults to a 64px border.
         The full image is 1024x1024.
         """
-        pixels = 0 if pixels < 0 else 512 if pixels > 512 else pixels
-
-        option = option.lower()
-
-        if option not in OPTIONS.keys():
+        option, pixels, flag = self.process_options(option, pixels)
+        if flag is None:
             return await ctx.send("I don't have that flag!")
 
-        flag = OPTIONS[option]
-
         async with ctx.typing():
-
-            # Get avatar bytes
             image_bytes = await ctx.author.avatar_url.read()
-            avatar = Image.open(BytesIO(image_bytes))
-            avatar = avatar.convert("RGBA").resize((1024, 1024))
-
-            avatar = self.crop_avatar(avatar)
-
-            ring = Image.open(Path(f"bot/resources/pride/flags/{flag}.png")).resize((1024, 1024))
-            ring = ring.convert("RGBA")
-            ring = self.crop_ring(ring, pixels)
+            await self.process_image(ctx, image_bytes, pixels, flag, option)
 
-            avatar.alpha_composite(ring, (0, 0))
-            bufferedio = BytesIO()
-            avatar.save(bufferedio, format="PNG")
-            bufferedio.seek(0)
+    @prideavatar.command()
+    async def image(self, ctx: Context, url: str, option: str = "lgbt", pixels: int = 64) -> None:
+        """
+        This surrounds the image specified by the URL with a border of a specified LGBT flag.
 
-            file = discord.File(bufferedio, filename="pride_avatar.png")  # Creates file to be used in embed
-            embed = discord.Embed(
-                name="Your Lovely Pride Avatar",
-                description=f"Here is your lovely avatar, surrounded by\n a beautiful {option} flag. Enjoy :D"
-            )
-            embed.set_image(url="attachment://pride_avatar.png")
-            embed.set_footer(text=f"Made by {ctx.author.display_name}", icon_url=ctx.author.avatar_url)
+        This defaults to the LGBT rainbow flag if none is given.
+        The amount of pixels can be given which determines the thickness of the flag border.
+        This has a maximum of 512px and defaults to a 64px border.
+        The full image is 1024x1024.
+        """
+        option, pixels, flag = self.process_options(option, pixels)
+        if flag is None:
+            return await ctx.send("I don't have that flag!")
 
-        await ctx.send(file=file, embed=embed)
+        async with ctx.typing():
+            async with aiohttp.ClientSession() as session:
+                try:
+                    response = await session.get(url)
+                except aiohttp.client_exceptions.ClientConnectorError:
+                    return await ctx.send("Cannot connect to provided URL!")
+                except aiohttp.client_exceptions.InvalidURL:
+                    return await ctx.send("Invalid URL!")
+                if response.status != 200:
+                    return await ctx.send("Bad response from provided URL!")
+                image_bytes = await response.read()
+                await self.process_image(ctx, image_bytes, pixels, flag, option)
 
     @prideavatar.command()
-    async def flags(self, ctx: commands.Context) -> None:
+    async def flags(self, ctx: Context) -> None:
         """This lists the flags that can be used with the prideavatar command."""
         choices = sorted(set(OPTIONS.values()))
         options = "• " + "\n• ".join(choices)
@@ -139,6 +172,6 @@ class PrideAvatar(commands.Cog):
         await ctx.send(embed=embed)
 
 
-def setup(bot: commands.Bot) -> None:
+def setup(bot: Bot) -> None:
     """Cog load."""
     bot.add_cog(PrideAvatar(bot))