mirror of
https://github.com/brandons209/Red-bot-Cogs.git
synced 2024-10-01 17:46:28 +13:00
226 lines
7.4 KiB
Python
226 lines
7.4 KiB
Python
|
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 <distort type> <intensity (1-10) (optional)> <image link>`
|
||
|
"""
|
||
|
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
|