diff options
Diffstat (limited to 'bot/exts/avatar_modification/avatar_modify.py')
| -rw-r--r-- | bot/exts/avatar_modification/avatar_modify.py | 372 |
1 files changed, 372 insertions, 0 deletions
diff --git a/bot/exts/avatar_modification/avatar_modify.py b/bot/exts/avatar_modification/avatar_modify.py new file mode 100644 index 00000000..87eb05e6 --- /dev/null +++ b/bot/exts/avatar_modification/avatar_modify.py @@ -0,0 +1,372 @@ +import asyncio +import json +import logging +import math +import string +import unicodedata +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Callable, Optional, TypeVar, Union + +import discord +from discord.ext import commands + +from bot.bot import Bot +from bot.constants import Colours, Emojis +from bot.exts.avatar_modification._effects import PfpEffects +from bot.utils.extensions import invoke_help_command +from bot.utils.halloween import spookifications + +log = logging.getLogger(__name__) + +_EXECUTOR = ThreadPoolExecutor(10) + +FILENAME_STRING = "{effect}_{author}.png" + +MAX_SQUARES = 10_000 + +T = TypeVar("T") + +GENDER_OPTIONS = json.loads(Path("bot/resources/holidays/pride/gender_options.json").read_text("utf8")) + + +async def in_executor(func: Callable[..., T], *args) -> T: + """ + Runs the given synchronous function `func` in an executor. + + This is useful for running slow, blocking code within async + functions, so that they don't block the bot. + """ + log.trace(f"Running {func.__name__} in an executor.") + loop = asyncio.get_event_loop() + return await loop.run_in_executor(_EXECUTOR, func, *args) + + +def file_safe_name(effect: str, display_name: str) -> str: + """Returns a file safe filename based on the given effect and display name.""" + valid_filename_chars = f"-_. {string.ascii_letters}{string.digits}" + + file_name = FILENAME_STRING.format(effect=effect, author=display_name) + + # Replace spaces + file_name = file_name.replace(" ", "_") + + # Normalize unicode characters + cleaned_filename = unicodedata.normalize("NFKD", file_name).encode("ASCII", "ignore").decode() + + # Remove invalid filename characters + cleaned_filename = "".join(c for c in cleaned_filename if c in valid_filename_chars) + return cleaned_filename + + +class AvatarModify(commands.Cog): + """Various commands for users to apply affects to their own avatars.""" + + def __init__(self, bot: Bot): + self.bot = bot + + async def _fetch_user(self, user_id: int) -> Optional[discord.User]: + """ + Fetches a user and handles errors. + + This helper function is required as the member cache doesn't always have the most up to date + profile picture. This can lead to errors if the image is deleted from the Discord CDN. + fetch_member can't be used due to the avatar url being part of the user object, and + some weird caching that D.py does + """ + try: + user = await self.bot.fetch_user(user_id) + except discord.errors.NotFound: + log.debug(f"User {user_id} could not be found.") + return None + except discord.HTTPException: + log.exception(f"Exception while trying to retrieve user {user_id} from Discord.") + return None + + return user + + @commands.group(aliases=("avatar_mod", "pfp_mod", "avatarmod", "pfpmod")) + async def avatar_modify(self, ctx: commands.Context) -> None: + """Groups all of the pfp modifying commands to allow a single concurrency limit.""" + if not ctx.invoked_subcommand: + await invoke_help_command(ctx) + + @avatar_modify.command(name="8bitify", root_aliases=("8bitify",)) + async def eightbit_command(self, ctx: commands.Context) -> None: + """Pixelates your avatar and changes the palette to an 8bit one.""" + async with ctx.typing(): + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + + image_bytes = await user.display_avatar.replace(size=1024).read() + file_name = file_safe_name("eightbit_avatar", ctx.author.display_name) + + file = await in_executor( + PfpEffects.apply_effect, + image_bytes, + PfpEffects.eight_bitify_effect, + file_name + ) + + embed = discord.Embed( + title="Your 8-bit avatar", + description="Here is your avatar. I think it looks all cool and 'retro'." + ) + + embed.set_image(url=f"attachment://{file_name}") + embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url) + + await ctx.send(embed=embed, file=file) + + @avatar_modify.command(name="reverse", root_aliases=("reverse",)) + async def reverse(self, ctx: commands.Context, *, text: Optional[str]) -> None: + """ + Reverses the sent text. + + If no text is provided, the user's profile picture will be reversed. + """ + if text: + await ctx.send(f"> {text[::-1]}", allowed_mentions=discord.AllowedMentions.none()) + return + + async with ctx.typing(): + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + + image_bytes = await user.display_avatar.replace(size=1024).read() + filename = file_safe_name("reverse_avatar", ctx.author.display_name) + + file = await in_executor( + PfpEffects.apply_effect, + image_bytes, + PfpEffects.flip_effect, + filename + ) + + embed = discord.Embed( + title="Your reversed avatar.", + description="Here is your reversed avatar. I think it is a spitting image of you." + ) + + embed.set_image(url=f"attachment://{filename}") + embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url) + + await ctx.send(embed=embed, file=file) + + @avatar_modify.command(aliases=("easterify",), root_aliases=("easterify", "avatareasterify")) + async def avatareasterify(self, ctx: commands.Context, *colours: Union[discord.Colour, str]) -> None: + """ + This "Easterifies" the user's avatar. + + Given colours will produce a personalised egg in the corner, similar to the egg_decorate command. + If colours are not given, a nice little chocolate bunny will sit in the corner. + Colours are split by spaces, unless you wrap the colour name in double quotes. + Discord colour names, HTML colour names, XKCD colour names and hex values are accepted. + """ + async def send(*args, **kwargs) -> str: + """ + This replaces the original ctx.send. + + When invoking the egg decorating command, the egg itself doesn't print to to the channel. + Returns the message content so that if any errors occur, the error message can be output. + """ + if args: + return args[0] + + async with ctx.typing(): + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + + egg = None + if colours: + send_message = ctx.send + ctx.send = send # Assigns ctx.send to a fake send + egg = await ctx.invoke(self.bot.get_command("eggdecorate"), *colours) + if isinstance(egg, str): # When an error message occurs in eggdecorate. + await send_message(egg) + return + ctx.send = send_message # Reassigns ctx.send + + image_bytes = await user.display_avatar.replace(size=256).read() + file_name = file_safe_name("easterified_avatar", ctx.author.display_name) + + file = await in_executor( + PfpEffects.apply_effect, + image_bytes, + PfpEffects.easterify_effect, + file_name, + egg + ) + + embed = discord.Embed( + title="Your Lovely Easterified Avatar!", + description="Here is your lovely avatar, all bright and colourful\nwith Easter pastel colours. Enjoy :D" + ) + embed.set_image(url=f"attachment://{file_name}") + embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url) + + await ctx.send(file=file, embed=embed) + + @staticmethod + async def send_pride_image( + ctx: commands.Context, + image_bytes: bytes, + pixels: int, + flag: str, + option: str + ) -> None: + """Gets and sends the image in an embed. Used by the pride commands.""" + async with ctx.typing(): + file_name = file_safe_name("pride_avatar", ctx.author.display_name) + + file = await in_executor( + PfpEffects.apply_effect, + image_bytes, + PfpEffects.pridify_effect, + file_name, + pixels, + flag + ) + + embed = discord.Embed( + title="Your Lovely Pride Avatar!", + description=f"Here is your lovely avatar, surrounded by\n a beautiful {option} flag. Enjoy :D" + ) + embed.set_image(url=f"attachment://{file_name}") + embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=ctx.author.avatar.url) + await ctx.send(file=file, embed=embed) + + @avatar_modify.group( + aliases=("avatarpride", "pridepfp", "prideprofile"), + root_aliases=("prideavatar", "avatarpride", "pridepfp", "prideprofile"), + invoke_without_command=True + ) + async def prideavatar(self, ctx: commands.Context, option: str = "lgbt", pixels: int = 64) -> None: + """ + This surrounds an avatar with a border of a specified LGBT flag. + + 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 = option.lower() + pixels = max(0, min(512, pixels)) + flag = GENDER_OPTIONS.get(option) + if flag is None: + await ctx.send("I don't have that flag!") + return + + async with ctx.typing(): + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + image_bytes = await user.display_avatar.replace(size=1024).read() + await self.send_pride_image(ctx, image_bytes, pixels, flag, option) + + @prideavatar.command() + async def flags(self, ctx: commands.Context) -> None: + """This lists the flags that can be used with the prideavatar command.""" + choices = sorted(set(GENDER_OPTIONS.values())) + options = "• " + "\n• ".join(choices) + embed = discord.Embed( + title="I have the following flags:", + description=options, + colour=Colours.soft_red + ) + await ctx.send(embed=embed) + + @avatar_modify.command( + aliases=("savatar", "spookify"), + root_aliases=("spookyavatar", "spookify", "savatar"), + brief="Spookify an user's avatar." + ) + async def spookyavatar(self, ctx: commands.Context) -> None: + """This "spookifies" the user's avatar, with a random *spooky* effect.""" + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + + async with ctx.typing(): + image_bytes = await user.display_avatar.replace(size=1024).read() + + file_name = file_safe_name("spooky_avatar", ctx.author.display_name) + + file = await in_executor( + PfpEffects.apply_effect, + image_bytes, + spookifications.get_random_effect, + file_name + ) + + embed = discord.Embed( + title="Is this you or am I just really paranoid?", + colour=Colours.soft_red + ) + embed.set_image(url=f"attachment://{file_name}") + embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=ctx.author.display_avatar.url) + + await ctx.send(file=file, embed=embed) + + @avatar_modify.command(name="mosaic", root_aliases=("mosaic",)) + async def mosaic_command(self, ctx: commands.Context, squares: int = 16) -> None: + """Splits your avatar into x squares, randomizes them and stitches them back into a new image!""" + async with ctx.typing(): + user = await self._fetch_user(ctx.author.id) + if not user: + await ctx.send(f"{Emojis.cross_mark} Could not get user info.") + return + + if not 1 <= squares <= MAX_SQUARES: + raise commands.BadArgument(f"Squares must be a positive number less than or equal to {MAX_SQUARES:,}.") + + sqrt = math.sqrt(squares) + + if not sqrt.is_integer(): + squares = math.ceil(sqrt) ** 2 # Get the next perfect square + + file_name = file_safe_name("mosaic_avatar", ctx.author.display_name) + + img_bytes = await user.display_avatar.replace(size=1024).read() + + file = await in_executor( + PfpEffects.apply_effect, + img_bytes, + PfpEffects.mosaic_effect, + file_name, + squares, + ) + + if squares == 1: + title = "Hooh... that was a lot of work" + description = "I present to you... Yourself!" + elif squares == MAX_SQUARES: + title = "Testing the limits I see..." + description = "What a masterpiece. :star:" + else: + title = "Your mosaic avatar" + description = f"Here is your avatar. I think it looks a bit *puzzling*\nMade with {squares} squares." + + embed = discord.Embed( + title=title, + description=description, + colour=Colours.blue + ) + + embed.set_image(url=f"attachment://{file_name}") + embed.set_footer(text=f"Made by {ctx.author.display_name}", icon_url=user.display_avatar.url) + + await ctx.send(file=file, embed=embed) + + +def setup(bot: Bot) -> None: + """Load the AvatarModify cog.""" + bot.add_cog(AvatarModify(bot)) |