Split effects from cog

3 files changed, 133 insertions, 125 deletions

diff --git a/bot/exts/evergreen/profile_pic_modification/__init__.py b/bot/exts/evergreen/profile_pic_modification/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/bot/exts/evergreen/profile_pic_modification/__init__.py
diff --git a/bot/exts/evergreen/profile_pic_modification/_effects.py b/bot/exts/evergreen/profile_pic_modification/_effects.py
new file mode 100644
index 00000000..fbe5f706
--- /dev/null
+++ b/bot/exts/evergreen/profile_pic_modification/_effects.py
@@ -0,0 +1,122 @@
+import typing as t
+from io import BytesIO
+from pathlib import Path
+
+import discord
+from PIL import Image, ImageDraw, ImageOps
+
+EASTER_COLOURS = [
+    (255, 247, 0), (255, 255, 224), (0, 255, 127), (189, 252, 201), (255, 192, 203),
+    (255, 160, 122), (181, 115, 220), (221, 160, 221), (200, 162, 200), (238, 130, 238),
+    (135, 206, 235), (0, 204, 204), (64, 224, 208)
+]  # Pastel colours - Easter-like
+
+
+class PfpEffects():
+    """Implements various image effects."""
+
+    @staticmethod
+    def closest(x: t.Tuple[int, int, int]) -> t.Tuple[int, int, int]:
+        """
+        Finds the closest easter colour to a given pixel.
+
+        Returns a merge between the original colour and the closest colour
+        """
+        r1, g1, b1 = x
+
+        def distance(point: t.Tuple[int, int, int]) -> t.Tuple[int, int, int]:
+            """Finds the difference between a pastel colour and the original pixel colour."""
+            r2, g2, b2 = point
+            return ((r1 - r2)**2 + (g1 - g2)**2 + (b1 - b2)**2)
+
+        closest_colours = sorted(EASTER_COLOURS, key=lambda point: distance(point))
+        r2, g2, b2 = closest_colours[0]
+        r = (r1 + r2) // 2
+        g = (g1 + g2) // 2
+        b = (b1 + b2) // 2
+
+        return (r, g, b)
+
+    @staticmethod
+    def crop_avatar_circle(avatar: Image) -> Image:
+        """This crops the avatar given into a circle."""
+        mask = Image.new("L", avatar.size, 0)
+        draw = ImageDraw.Draw(mask)
+        draw.ellipse((0, 0) + avatar.size, fill=255)
+        avatar.putalpha(mask)
+        return avatar
+
+    @staticmethod
+    def crop_ring(ring: Image, px: int) -> Image:
+        """This crops the given ring into a circle."""
+        mask = Image.new("L", ring.size, 0)
+        draw = ImageDraw.Draw(mask)
+        draw.ellipse((0, 0) + ring.size, fill=255)
+        draw.ellipse((px, px, 1024-px, 1024-px), fill=0)
+        ring.putalpha(mask)
+        return ring
+
+    @staticmethod
+    def _apply_effect(image_bytes: bytes, effect: t.Callable, *args) -> discord.File:
+        im = Image.open(BytesIO(image_bytes))
+        im = im.convert("RGBA")
+        im = effect(im, *args)
+
+        bufferedio = BytesIO()
+        im.save(bufferedio, format="PNG")
+        bufferedio.seek(0)
+
+        return discord.File(bufferedio, filename="modified_avatar.png")
+
+    @staticmethod
+    def pridify_effect(
+        image: Image,
+        pixels: int,
+        flag: str
+    ) -> Image:
+        """Applies the pride effect to the given image."""
+        image = image.resize((1024, 1024))
+        image = PfpEffects.crop_avatar_circle(image)
+
+        ring = Image.open(Path(f"bot/resources/pride/flags/{flag}.png")).resize((1024, 1024))
+        ring = ring.convert("RGBA")
+        ring = PfpEffects.crop_ring(ring, pixels)
+
+        image.alpha_composite(ring, (0, 0))
+        return image
+
+    @staticmethod
+    def eight_bitify_effect(image: Image) -> Image:
+        """Applies the 8bit effect to the given image."""
+        image = image.resize((32, 32), resample=Image.NEAREST).resize((1024, 1024), resample=Image.NEAREST)
+        return image.quantize()
+
+    @staticmethod
+    def easterify_effect(image: Image, overlay_image: Image = None) -> Image:
+        """Applies the easter effect to the given image."""
+        if overlay_image:
+            ratio = 64 / overlay_image.height
+            overlay_image = overlay_image.resize((
+                round(overlay_image.width * ratio),
+                round(overlay_image.height * ratio)
+            ))
+            overlay_image = overlay_image.convert("RGBA")
+        else:
+            overlay_image = overlay_image = Image.open(Path("bot/resources/easter/chocolate_bunny.png"))
+
+        alpha = image.getchannel("A").getdata()
+        image = image.convert("RGB")
+        image = ImageOps.posterize(image, 6)
+
+        data = image.getdata()
+        setted_data = set(data)
+        new_d = {}
+
+        for x in setted_data:
+            new_d[x] = PfpEffects.closest(x)
+        new_data = [(*new_d[x], alpha[i]) if x in new_d else x for i, x in enumerate(data)]
+
+        im = Image.new("RGBA", image.size)
+        im.putdata(new_data)
+        im.alpha_composite(overlay_image, (im.width - overlay_image.width, (im.height - overlay_image.height)//2))
+        return im
diff --git a/bot/exts/evergreen/pfp_modify.py b/bot/exts/evergreen/profile_pic_modification/pfp_modify.py
index 5fb3b89d..51742257 100644
--- a/bot/exts/evergreen/pfp_modify.py
+++ b/bot/exts/evergreen/profile_pic_modification/pfp_modify.py
@@ -3,26 +3,18 @@ import json
 import logging
 import typing as t
 from concurrent.futures import ThreadPoolExecutor
-from io import BytesIO
-from pathlib import Path
 
 import aiofiles
 import discord
-from PIL import Image, ImageDraw, ImageOps
 from aiohttp import client_exceptions
 from discord.ext import commands
 
 from bot.constants import Colours
+from bot.exts.evergreen.profile_pic_modification._effects import PfpEffects
 from bot.utils.halloween import spookifications
 
 log = logging.getLogger(__name__)
 
-EASTER_COLOURS = [
-    (255, 247, 0), (255, 255, 224), (0, 255, 127), (189, 252, 201), (255, 192, 203),
-    (255, 160, 122), (181, 115, 220), (221, 160, 221), (200, 162, 200), (238, 130, 238),
-    (135, 206, 235), (0, 204, 204), (64, 224, 208)
-]  # Pastel colours - Easter-like
-
 _EXECUTOR = ThreadPoolExecutor(10)
 
 
@@ -45,112 +37,6 @@ class PfpModify(commands.Cog):
         async with aiofiles.open('bot/resources/pride/gender_options.json') as f:
             self.GENDER_OPTIONS = json.loads(await f.read())
 
-    @staticmethod
-    def closest(x: t.Tuple[int, int, int]) -> t.Tuple[int, int, int]:
-        """
-        Finds the closest easter colour to a given pixel.
-
-        Returns a merge between the original colour and the closest colour
-        """
-        r1, g1, b1 = x
-
-        def distance(point: t.Tuple[int, int, int]) -> t.Tuple[int, int, int]:
-            """Finds the difference between a pastel colour and the original pixel colour."""
-            r2, g2, b2 = point
-            return ((r1 - r2)**2 + (g1 - g2)**2 + (b1 - b2)**2)
-
-        closest_colours = sorted(EASTER_COLOURS, key=lambda point: distance(point))
-        r2, g2, b2 = closest_colours[0]
-        r = (r1 + r2) // 2
-        g = (g1 + g2) // 2
-        b = (b1 + b2) // 2
-
-        return (r, g, b)
-
-    @staticmethod
-    def crop_avatar_circle(avatar: Image) -> Image:
-        """This crops the avatar given into a circle."""
-        mask = Image.new("L", avatar.size, 0)
-        draw = ImageDraw.Draw(mask)
-        draw.ellipse((0, 0) + avatar.size, fill=255)
-        avatar.putalpha(mask)
-        return avatar
-
-    @staticmethod
-    def crop_ring(ring: Image, px: int) -> Image:
-        """This crops the given ring into a circle."""
-        mask = Image.new("L", ring.size, 0)
-        draw = ImageDraw.Draw(mask)
-        draw.ellipse((0, 0) + ring.size, fill=255)
-        draw.ellipse((px, px, 1024-px, 1024-px), fill=0)
-        ring.putalpha(mask)
-        return ring
-
-    @staticmethod
-    def _apply_effect(image_bytes: bytes, effect: t.Callable, *args) -> discord.File:
-        im = Image.open(BytesIO(image_bytes))
-        im = im.convert("RGBA")
-        im = effect(im, *args)
-
-        bufferedio = BytesIO()
-        im.save(bufferedio, format="PNG")
-        bufferedio.seek(0)
-
-        return discord.File(bufferedio, filename="modified_avatar.png")
-
-    @staticmethod
-    def pridify_effect(
-        image: Image,
-        pixels: int,
-        flag: str
-    ) -> Image:
-        """Applies the pride effect to the given image."""
-        image = image.resize((1024, 1024))
-        image = PfpModify.crop_avatar_circle(image)
-
-        ring = Image.open(Path(f"bot/resources/pride/flags/{flag}.png")).resize((1024, 1024))
-        ring = ring.convert("RGBA")
-        ring = PfpModify.crop_ring(ring, pixels)
-
-        image.alpha_composite(ring, (0, 0))
-        return image
-
-    @staticmethod
-    def eight_bitify_effect(image: Image) -> Image:
-        """Applies the 8bit effect to the given image."""
-        image = image.resize((32, 32), resample=Image.NEAREST).resize((1024, 1024), resample=Image.NEAREST)
-        return image.quantize()
-
-    @staticmethod
-    def easterify_effect(image: Image, overlay_image: Image = None) -> Image:
-        """Applies the easter effect to the given image."""
-        if overlay_image:
-            ratio = 64 / overlay_image.height
-            overlay_image = overlay_image.resize((
-                round(overlay_image.width * ratio),
-                round(overlay_image.height * ratio)
-            ))
-            overlay_image = overlay_image.convert("RGBA")
-        else:
-            overlay_image = overlay_image = Image.open(Path("bot/resources/easter/chocolate_bunny.png"))
-
-        alpha = image.getchannel("A").getdata()
-        image = image.convert("RGB")
-        image = ImageOps.posterize(image, 6)
-
-        data = image.getdata()
-        setted_data = set(data)
-        new_d = {}
-
-        for x in setted_data:
-            new_d[x] = PfpModify.closest(x)
-        new_data = [(*new_d[x], alpha[i]) if x in new_d else x for i, x in enumerate(data)]
-
-        im = Image.new("RGBA", image.size)
-        im.putdata(new_data)
-        im.alpha_composite(overlay_image, (im.width - overlay_image.width, (im.height - overlay_image.height)//2))
-        return im
-
     @commands.group()
     async def pfp_modify(self, ctx: commands.Context) -> None:
         """Groups all of the pfp modifing commands to allow a single concurrency limit."""
@@ -163,9 +49,9 @@ class PfpModify(commands.Cog):
         async with ctx.typing():
             image_bytes = await ctx.author.avatar_url.read()
             file = await in_thread(
-                self._apply_effect,
+                PfpEffects._apply_effect,
                 image_bytes,
-                self.eight_bitify_effect
+                PfpEffects.eight_bitify_effect
             )
 
             embed = discord.Embed(
@@ -211,9 +97,9 @@ class PfpModify(commands.Cog):
 
             image_bytes = await ctx.author.avatar_url_as(size=256).read()
             file = await in_thread(
-                self._apply_effect,
+                PfpEffects._apply_effect,
                 image_bytes,
-                self.easterify_effect,
+                PfpEffects.easterify_effect,
                 egg
             )
 
@@ -237,9 +123,9 @@ class PfpModify(commands.Cog):
         """Gets and sends the image in an embed. Used by the pride commands."""
         async with ctx.typing():
             file = await in_thread(
-                self._apply_effect,
+                PfpEffects._apply_effect,
                 image_bytes,
-                self.pridify_effect,
+                PfpEffects.pridify_effect,
                 pixels,
                 flag
             )
@@ -318,9 +204,9 @@ class PfpModify(commands.Cog):
         await ctx.send(embed=embed)
 
     @pfp_modify.command(
-        name='savatar',
-        aliases=('spookyavatar', 'spookify'),
-        root_aliases=('spookyavatar', 'spookify'),
+        name='spookyavatar',
+        aliases=('savatar', 'spookify'),
+        root_aliases=('spookyavatar', 'spookify', 'savatar'),
         brief='Spookify an user\'s avatar.'
     )
     async def spooky_avatar(self, ctx: commands.Context, user: discord.Member = None) -> None:
@@ -330,7 +216,7 @@ class PfpModify(commands.Cog):
 
         async with ctx.typing():
             image_bytes = await ctx.author.avatar_url.read()
-            file = await in_thread(self._apply_effect, image_bytes, spookifications.get_random_effect)
+            file = await in_thread(PfpEffects._apply_effect, image_bytes, spookifications.get_random_effect)
 
             embed = discord.Embed(
                 title="Is this you or am I just really paranoid?",