mirror of
https://github.com/uhIgnacio/EmoteManager.git
synced 2024-08-15 02:23:13 +00:00
auto convert static images to GIFs if there's no room (closes #3)
Also backport the image resize code from 70045b2a0e
.
This commit is contained in:
parent
1ad3a7ccec
commit
e02022b245
4 changed files with 101 additions and 50 deletions
4
.gitignore
vendored
4
.gitignore
vendored
|
@ -8,3 +8,7 @@ __pycache__/
|
|||
venv/
|
||||
|
||||
config.py
|
||||
|
||||
*.png
|
||||
*.gif
|
||||
*.jpg
|
||||
|
|
|
@ -3,17 +3,20 @@
|
|||
|
||||
import asyncio
|
||||
import cgi
|
||||
import collections
|
||||
import contextlib
|
||||
import io
|
||||
import logging
|
||||
import weakref
|
||||
import operator
|
||||
import posixpath
|
||||
import traceback
|
||||
import contextlib
|
||||
import urllib.parse
|
||||
import weakref
|
||||
|
||||
import aioec
|
||||
import aiohttp
|
||||
import discord
|
||||
import humanize
|
||||
from discord.ext import commands
|
||||
|
||||
import utils
|
||||
|
@ -24,6 +27,9 @@ from utils.paginator import ListPaginator
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class UserCancelledError(commands.UserInputError):
|
||||
pass
|
||||
|
||||
class Emotes(commands.Cog):
|
||||
IMAGE_MIMETYPES = {'image/png', 'image/jpeg', 'image/gif'}
|
||||
# TAR_MIMETYPES = {'application/x-tar', 'application/x-xz', 'application/gzip', 'application/x-bzip2'}
|
||||
|
@ -107,7 +113,7 @@ class Emotes(commands.Cog):
|
|||
"""
|
||||
name, url = self.parse_add_command_args(context, args)
|
||||
async with context.typing():
|
||||
message = await self.add_safe(context.guild, name, url, context.message.author.id)
|
||||
message = await self.add_safe(context, name, url, context.message.author.id)
|
||||
await context.send(message)
|
||||
|
||||
@classmethod
|
||||
|
@ -182,7 +188,7 @@ class Emotes(commands.Cog):
|
|||
f'Original emote author: {utils.format_user(self.bot, emote.author)}')
|
||||
|
||||
async with context.typing():
|
||||
message = await self.add_safe(context.guild, name, emote.url, context.author.id, reason=reason)
|
||||
message = await self.add_safe(context, name, emote.url, context.author.id, reason=reason)
|
||||
|
||||
await context.send(message)
|
||||
|
||||
|
@ -199,12 +205,16 @@ class Emotes(commands.Cog):
|
|||
raise commands.BadArgument('A URL or attachment must be given.')
|
||||
|
||||
url = url or context.message.attachments[0].url
|
||||
async with context.typing():
|
||||
archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES)
|
||||
if type(archive) is str: # error case
|
||||
await context.send(archive)
|
||||
return
|
||||
|
||||
await self.add_from_archive(context, archive)
|
||||
with contextlib.suppress(discord.HTTPException):
|
||||
# so they know when we're done
|
||||
await context.message.add_reaction(utils.SUCCESS_EMOJIS[True])
|
||||
|
||||
async def add_from_archive(self, context, archive):
|
||||
limit = 50_000_000 # prevent someone from trying to make a giant compressed file
|
||||
|
@ -212,7 +222,8 @@ class Emotes(commands.Cog):
|
|||
if error is None:
|
||||
name = self.format_emote_filename(posixpath.basename(name))
|
||||
async with context.typing():
|
||||
await context.send(await self.add_safe_bytes(context.guild, name, context.author.id, img))
|
||||
message = await self.add_safe_bytes(context, name, context.author.id, img)
|
||||
await context.send(message)
|
||||
continue
|
||||
|
||||
if isinstance(error, errors.FileTooBigError):
|
||||
|
@ -224,7 +235,7 @@ class Emotes(commands.Cog):
|
|||
|
||||
await context.send(f'{name}: {error}')
|
||||
|
||||
async def add_safe(self, guild, name, url, author_id, *, reason=None):
|
||||
async def add_safe(self, context, name, url, author_id, *, reason=None):
|
||||
"""Try to add an emote. Returns a string that should be sent to the user."""
|
||||
try:
|
||||
image_data = await self.fetch_safe(url)
|
||||
|
@ -233,7 +244,7 @@ class Emotes(commands.Cog):
|
|||
|
||||
if type(image_data) is str: # error case
|
||||
return image_data
|
||||
return await self.add_safe_bytes(guild, name, author_id, image_data, reason=reason)
|
||||
return await self.add_safe_bytes(context, name, author_id, image_data, reason=reason)
|
||||
|
||||
async def fetch_safe(self, url, valid_mimetypes=None):
|
||||
"""Try to fetch a URL. On error return a string that should be sent to the user."""
|
||||
|
@ -246,17 +257,34 @@ class Emotes(commands.Cog):
|
|||
except aiohttp.ClientResponseError as exc:
|
||||
raise errors.HTTPException(exc.status)
|
||||
|
||||
async def add_safe_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None):
|
||||
"""Try to add an emote from bytes. On error, return a string that should be sent to the user."""
|
||||
async def add_safe_bytes(self, context, name, author_id, image_data: bytes, *, reason=None):
|
||||
"""Try to add an emote from bytes. On error, return a string that should be sent to the user.
|
||||
|
||||
If the image is static and there are not enough free static slots, prompt to convert the image to a single-frame
|
||||
gif instead.
|
||||
"""
|
||||
counts = collections.Counter(map(operator.attrgetter('animated'), context.guild.emojis))
|
||||
# >= rather than == because there are sneaky ways to exceed the limit
|
||||
if counts[False] >= context.guild.emoji_limit and counts[True] >= context.guild.emoji_limit:
|
||||
# we raise instead of returning a string in order to abort commands that run this function in a loop
|
||||
raise commands.UserInputError('This server is out of emote slots.')
|
||||
|
||||
static = utils.image.mime_type_for_image(image_data) != 'image/gif'
|
||||
converted = False
|
||||
if static and counts[False] >= 1: # context.guild.emoji_limit:
|
||||
image_data = await utils.image.convert_to_gif_in_subprocess(image_data)
|
||||
converted = True
|
||||
|
||||
try:
|
||||
emote = await self.create_emote_from_bytes(guild, name, author_id, image_data, reason=reason)
|
||||
emote = await self.create_emote_from_bytes(context.guild, name, author_id, image_data, reason=reason)
|
||||
except discord.InvalidArgument:
|
||||
return discord.utils.escape_mentions(f'{name}: The file supplied was not a valid GIF, PNG, or JPEG file.')
|
||||
except discord.HTTPException as ex:
|
||||
return discord.utils.escape_mentions(
|
||||
f'{name}: An error occurred while creating the the emote:\n'
|
||||
+ utils.format_http_exception(ex))
|
||||
return f'Emote {emote} successfully created.'
|
||||
s = f'Emote {emote} successfully created'
|
||||
return s + ' as a GIF.' if converted else s + '.'
|
||||
|
||||
async def fetch(self, url, valid_mimetypes=None):
|
||||
valid_mimetypes = valid_mimetypes or self.IMAGE_MIMETYPES
|
||||
|
|
|
@ -17,11 +17,18 @@ class EmoteManagerError(commands.CommandError):
|
|||
"""Generic error with the bot. This can be used to catch all bot errors."""
|
||||
pass
|
||||
|
||||
class ImageResizeTimeoutError(EmoteManagerError, asyncio.TimeoutError):
|
||||
class ImageProcessingTimeoutError(EmoteManagerError, asyncio.TimeoutError):
|
||||
pass
|
||||
|
||||
class ImageResizeTimeoutError(ImageProcessingTimeoutError):
|
||||
"""Resizing the image took too long."""
|
||||
def __init__(self):
|
||||
super().__init__('Error: resizing the image took too long.')
|
||||
|
||||
class ImageConversionTimeoutError(ImageProcessingTimeoutError):
|
||||
def __init__(self):
|
||||
super().__init__('Error: converting the image to a GIF took too long.')
|
||||
|
||||
class HTTPException(EmoteManagerError):
|
||||
"""The server did not respond with an OK status code."""
|
||||
def __init__(self, status):
|
||||
|
|
|
@ -4,6 +4,7 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import functools
|
||||
import io
|
||||
import logging
|
||||
import sys
|
||||
|
@ -28,42 +29,43 @@ def resize_until_small(image_data: io.BytesIO) -> None:
|
|||
# so resizing sometimes does more harm than good.
|
||||
max_resolution = 128 # pixels
|
||||
image_size = size(image_data)
|
||||
while image_size > 256 * 2**10 and max_resolution >= 32: # don't resize past 32x32 or 256KiB
|
||||
if image_size <= 256 * 2**10:
|
||||
return
|
||||
|
||||
try:
|
||||
with wand.image.Image(blob=image_data) as original_image:
|
||||
while True:
|
||||
logger.debug('image size too big (%s bytes)', image_size)
|
||||
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
|
||||
|
||||
try:
|
||||
thumbnail(image_data, (max_resolution, max_resolution))
|
||||
with original_image.clone() as resized:
|
||||
resized.transform(resize=f'{max_resolution}x{max_resolution}')
|
||||
image_size = len(resized.make_blob())
|
||||
if image_size <= 256 * 2**10 or max_resolution < 32: # don't resize past 256KiB or 32×32
|
||||
image_data.truncate(0)
|
||||
image_data.seek(0)
|
||||
resized.save(file=image_data)
|
||||
image_data.seek(0)
|
||||
break
|
||||
|
||||
max_resolution //= 2
|
||||
except wand.exceptions.CoderError:
|
||||
raise errors.InvalidImageError
|
||||
|
||||
image_size = size(image_data)
|
||||
max_resolution //= 2
|
||||
def convert_to_gif(image_data: io.BytesIO) -> None:
|
||||
try:
|
||||
with wand.image.Image(blob=image_data) as orig, orig.convert('gif') as converted:
|
||||
# discord tries to stop us from abusing animated gif slots by detecting single frame gifs
|
||||
# so make it two frames
|
||||
converted.sequence[0].delay = 0 # show the first frame forever
|
||||
converted.sequence.append(wand.image.Image(width=1, height=1))
|
||||
|
||||
def thumbnail(image_data: io.BytesIO, max_size=(128, 128)) -> None:
|
||||
"""Resize an image in place to no more than max_size pixels, preserving aspect ratio."""
|
||||
with wand.image.Image(blob=image_data) as image:
|
||||
new_resolution = scale_resolution((image.width, image.height), max_size)
|
||||
image.resize(*new_resolution)
|
||||
image_data.truncate(0)
|
||||
image_data.seek(0)
|
||||
image.save(file=image_data)
|
||||
|
||||
# allow resizing the original image more than once for memory profiling
|
||||
converted.save(file=image_data)
|
||||
image_data.seek(0)
|
||||
|
||||
def scale_resolution(old_res, new_res):
|
||||
"""Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res"""
|
||||
# https://stackoverflow.com/a/6565988
|
||||
|
||||
old_width, old_height = old_res
|
||||
new_width, new_height = new_res
|
||||
|
||||
old_ratio = old_width / old_height
|
||||
new_ratio = new_width / new_height
|
||||
if new_ratio > old_ratio:
|
||||
return (old_width * new_height//old_height, new_height)
|
||||
return new_width, old_height * new_width//old_width
|
||||
except wand.exceptions.CoderError:
|
||||
raise errors.InvalidImageError
|
||||
|
||||
def mime_type_for_image(data):
|
||||
if data.startswith(b'\x89PNG\r\n\x1a\n'):
|
||||
|
@ -81,12 +83,19 @@ def image_to_base64_url(data):
|
|||
return fmt.format(mime=mime, data=b64)
|
||||
|
||||
def main() -> typing.NoReturn:
|
||||
"""resize an image from stdin and write the resized version to stdout."""
|
||||
"""resize or convert an image from stdin and write the resized or converted version to stdout."""
|
||||
import sys
|
||||
|
||||
if sys.argv[1] == 'resize':
|
||||
f = resize_until_small
|
||||
elif sys.argv[1] == 'convert':
|
||||
f = convert_to_gif
|
||||
else:
|
||||
sys.exit(1)
|
||||
|
||||
data = io.BytesIO(sys.stdin.buffer.read())
|
||||
try:
|
||||
resize_until_small(data)
|
||||
f(data)
|
||||
except errors.InvalidImageError:
|
||||
# 2 is used because 1 is already used by python's default error handler
|
||||
sys.exit(2)
|
||||
|
@ -102,19 +111,19 @@ def main() -> typing.NoReturn:
|
|||
|
||||
sys.exit(0)
|
||||
|
||||
async def resize_in_subprocess(image_data: bytes):
|
||||
async def process_image_in_subprocess(command_name, image_data: bytes):
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
sys.executable, '-m', __name__,
|
||||
sys.executable, '-m', __name__, command_name,
|
||||
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE)
|
||||
|
||||
try:
|
||||
image_data, err = await asyncio.wait_for(proc.communicate(image_data), timeout=30)
|
||||
image_data, err = await asyncio.wait_for(proc.communicate(image_data), timeout=float('inf'))
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
raise errors.ImageResizeTimeoutError
|
||||
raise errors.ImageResizeTimeoutError if command_name == 'resize' else errors.ImageConversionTimeoutError
|
||||
else:
|
||||
if proc.returncode == 2:
|
||||
raise errors.InvalidImageError
|
||||
|
@ -123,6 +132,9 @@ async def resize_in_subprocess(image_data: bytes):
|
|||
|
||||
return image_data
|
||||
|
||||
resize_in_subprocess = functools.partial(process_image_in_subprocess, 'resize')
|
||||
convert_to_gif_in_subprocess = functools.partial(process_image_in_subprocess, 'convert')
|
||||
|
||||
def size(fp):
|
||||
"""return the size, in bytes, of the data a file-like object represents"""
|
||||
with preserve_position(fp):
|
||||
|
|
Loading…
Reference in a new issue