import aiohttp, discord from redbot.core import Config, commands from wand.image import Image from io import BytesIO from typing import Optional, Tuple import asyncio, functools, urllib MAX_SIZE = 8 * 1024 * 1024 # by Flame442 class ImageFindError(Exception): """Generic error for the __get_image function.""" pass class ImageMagic(commands.Cog): def __init__(self, bot): super().__init__() self.config = Config.get_conf(self, identifier=4928034571, force_registration=True) self.bot = bot async def _get_image(self, ctx, link: str = None) -> Image: if ctx.guild: max_filesize = ctx.guild.filesize_limit else: max_filesize = MAX_SIZE # original by Flame442, edited for Wand by ScriptPony if not ctx.message.attachments and not link: async for msg in ctx.channel.history(limit=10): for a in msg.attachments: path = urllib.parse.urlparse(a.url).path link = a.url break if link: break if not link: raise ImageFindError("Please provide an attachment.") if link: # linked image path = urllib.parse.urlparse(link).path async with aiohttp.ClientSession() as session: try: async with session.get(link) as response: r = await response.read() try: img = Image(file=BytesIO(r)) except: raise ImageFindError("Invalid filetype") except (OSError, aiohttp.ClientError): raise ImageFindError("An image could not be found. Make sure you provide a direct link.") else: # attached image path = urllib.parse.urlparse(ctx.message.attachments[0].url).path if ctx.message.attachments[0].size > max_filesize: raise ImageFindError("That image is too large.") temp_orig = BytesIO() await ctx.message.attachments[0].save(temp_orig) temp_orig.seek(0) try: img = Image(file=temp_orig) except: raise ImageFindError("Invalid filetype") return img @staticmethod def _intensity(intensity: float) -> float: if intensity < 0: intensity = 0 elif intensity > 10: intensity = 10 intensity /= 10 return intensity def _distortion(self, img: Image, func: str, args: Tuple) -> Tuple[Image, str]: # distort img.iterator_reset() function = getattr(img, func, None) if function is None: return function(*args) if img.animation: while img.iterator_next(): function(*args) # image object and filename return img, (f"{func}." + img.mimetype[img.mimetype.find("/") + 1 :]) async def _command_body(self, ctx, args: Tuple): task = functools.partial(*args) task = self.bot.loop.run_in_executor(None, task) try: img, name = await asyncio.wait_for(task, timeout=60) except asyncio.TimeoutError: await ctx.send("The image took too long to process.") return try: await ctx.send(file=discord.File(BytesIO(img.make_blob()), name)) except discord.errors.HTTPException: await ctx.send("That image is too large.") return @commands.group() @commands.bot_has_permissions(attach_files=True) async def distort(self, ctx): """ Distorts an image from a direct link, attatchment, or from recent chat messages `[p]distort ` """ pass @distort.command() async def barrel(self, ctx, intensity: Optional[float] = 10, *, link: str = None): """ Bulges the center of the image outward """ intensity = self._intensity(intensity) amount = 0.3 async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) await self._command_body( ctx, args=( self._distortion, img, "distort", ("barrel", (amount * intensity, amount * intensity, amount * intensity, 0)), ), ) @distort.command() async def implode(self, ctx, intensity: Optional[float] = 10, *, link: str = None): """ Pinches in the center of the image """ intensity = self._intensity(intensity) amount = 0.6 async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) await self._command_body(ctx, args=(self._distortion, img, "implode", (amount * intensity,))) @distort.command() async def swirl(self, ctx, intensity: Optional[float] = 10, *, link: str = None): """ Swirls the center of the image """ switch = {0: 0, 1: 18, 2: 36, 3: 54, 4: 72, 5: 90, 6: 108, 7: 126, 8: 144, 9: 162, 10: 180} intensity = float(switch.get(round(intensity), 180)) async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) await self._command_body(ctx, args=(self._distortion, img, "swirl", (intensity,))) @distort.command() async def charcoal(self, ctx, intensity: Optional[float], *, link: str = None): """ Makes the image look somewhat like it was drawn with charcoal """ async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) await self._command_body(ctx, args=(self._distortion, img, "charcoal", (1.5, 0.5))) @distort.command() async def sketch(self, ctx, intensity: Optional[float], *, link: str = None): """ Makes the image look like it is a sketch """ async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) await self._command_body(ctx, args=(self._distortion, img, "sketch", (0.5, 0.0, 98.0))) @distort.command() async def zoom(self, ctx, intensity: Optional[float], *, link: str = None): """ Zooms in on the center of an image """ async with ctx.typing(): try: img = await self._get_image(ctx, link) except ImageFindError as e: return await ctx.send(e) h = img.height w = img.width img = self._distortion(img, "transform", (f"{w}x{h}", "150%"))[0] await self._command_body(ctx, args=(self._distortion, img, "transform", (f"{w/1.5}x{h/1.5}+{w/2}+{h/2}",))) async def red_delete_data_for_user( self, *, requester: Literal["discord_deleted_user", "owner", "user", "user_strict"], user_id: int, ): pass