From e02022b245872a48146145bcdbeb509c52ccf22a Mon Sep 17 00:00:00 2001 From: Io Mintz Date: Thu, 10 Oct 2019 00:24:12 +0000 Subject: [PATCH] auto convert static images to GIFs if there's no room (closes #3) Also backport the image resize code from https://github.com/EmoteCollector/EmoteCollector/commit/70045b2a0ec49e661062e271b773e52c22f8a9f6. --- .gitignore | 4 +++ cogs/emote.py | 52 +++++++++++++++++++++++------- utils/errors.py | 9 +++++- utils/image.py | 86 ++++++++++++++++++++++++++++--------------------- 4 files changed, 101 insertions(+), 50 deletions(-) diff --git a/.gitignore b/.gitignore index 1cd8fc2..0516f38 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,7 @@ __pycache__/ venv/ config.py + +*.png +*.gif +*.jpg diff --git a/cogs/emote.py b/cogs/emote.py index e1703fe..30a738f 100644 --- a/cogs/emote.py +++ b/cogs/emote.py @@ -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 - archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES) + 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 diff --git a/utils/errors.py b/utils/errors.py index d9bf44d..1cb5236 100644 --- a/utils/errors.py +++ b/utils/errors.py @@ -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): diff --git a/utils/image.py b/utils/image.py index 8d3ffbd..af845c5 100755 --- a/utils/image.py +++ b/utils/image.py @@ -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 - logger.debug('image size too big (%s bytes)', image_size) - logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution) + if image_size <= 256 * 2**10: + return - try: - thumbnail(image_data, (max_resolution, max_resolution)) - except wand.exceptions.CoderError: - raise errors.InvalidImageError + 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) - image_size = size(image_data) - max_resolution //= 2 + 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 -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) + max_resolution //= 2 + except wand.exceptions.CoderError: + raise errors.InvalidImageError - # allow resizing the original image more than once for memory profiling - image_data.seek(0) +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 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 + image_data.truncate(0) + image_data.seek(0) + converted.save(file=image_data) + image_data.seek(0) + 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):