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:
Io Mintz 2019-10-10 00:24:12 +00:00
parent 1ad3a7ccec
commit e02022b245
4 changed files with 101 additions and 50 deletions

4
.gitignore vendored
View File

@ -8,3 +8,7 @@ __pycache__/
venv/ venv/
config.py config.py
*.png
*.gif
*.jpg

View File

@ -3,17 +3,20 @@
import asyncio import asyncio
import cgi import cgi
import collections
import contextlib
import io import io
import logging import logging
import weakref import operator
import posixpath import posixpath
import traceback import traceback
import contextlib
import urllib.parse import urllib.parse
import weakref
import aioec import aioec
import aiohttp import aiohttp
import discord import discord
import humanize
from discord.ext import commands from discord.ext import commands
import utils import utils
@ -24,6 +27,9 @@ from utils.paginator import ListPaginator
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class UserCancelledError(commands.UserInputError):
pass
class Emotes(commands.Cog): class Emotes(commands.Cog):
IMAGE_MIMETYPES = {'image/png', 'image/jpeg', 'image/gif'} IMAGE_MIMETYPES = {'image/png', 'image/jpeg', 'image/gif'}
# TAR_MIMETYPES = {'application/x-tar', 'application/x-xz', 'application/gzip', 'application/x-bzip2'} # 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) name, url = self.parse_add_command_args(context, args)
async with context.typing(): 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) await context.send(message)
@classmethod @classmethod
@ -182,7 +188,7 @@ class Emotes(commands.Cog):
f'Original emote author: {utils.format_user(self.bot, emote.author)}') f'Original emote author: {utils.format_user(self.bot, emote.author)}')
async with context.typing(): 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) await context.send(message)
@ -199,12 +205,16 @@ class Emotes(commands.Cog):
raise commands.BadArgument('A URL or attachment must be given.') raise commands.BadArgument('A URL or attachment must be given.')
url = url or context.message.attachments[0].url 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 if type(archive) is str: # error case
await context.send(archive) await context.send(archive)
return return
await self.add_from_archive(context, archive) 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): async def add_from_archive(self, context, archive):
limit = 50_000_000 # prevent someone from trying to make a giant compressed file 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: if error is None:
name = self.format_emote_filename(posixpath.basename(name)) name = self.format_emote_filename(posixpath.basename(name))
async with context.typing(): 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 continue
if isinstance(error, errors.FileTooBigError): if isinstance(error, errors.FileTooBigError):
@ -224,7 +235,7 @@ class Emotes(commands.Cog):
await context.send(f'{name}: {error}') 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 to add an emote. Returns a string that should be sent to the user."""
try: try:
image_data = await self.fetch_safe(url) image_data = await self.fetch_safe(url)
@ -233,7 +244,7 @@ class Emotes(commands.Cog):
if type(image_data) is str: # error case if type(image_data) is str: # error case
return image_data 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): 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.""" """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: except aiohttp.ClientResponseError as exc:
raise errors.HTTPException(exc.status) raise errors.HTTPException(exc.status)
async def add_safe_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None): 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.""" """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: 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: except discord.InvalidArgument:
return discord.utils.escape_mentions(f'{name}: The file supplied was not a valid GIF, PNG, or JPEG file.') return discord.utils.escape_mentions(f'{name}: The file supplied was not a valid GIF, PNG, or JPEG file.')
except discord.HTTPException as ex: except discord.HTTPException as ex:
return discord.utils.escape_mentions( return discord.utils.escape_mentions(
f'{name}: An error occurred while creating the the emote:\n' f'{name}: An error occurred while creating the the emote:\n'
+ utils.format_http_exception(ex)) + 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): async def fetch(self, url, valid_mimetypes=None):
valid_mimetypes = valid_mimetypes or self.IMAGE_MIMETYPES valid_mimetypes = valid_mimetypes or self.IMAGE_MIMETYPES

View File

@ -17,11 +17,18 @@ class EmoteManagerError(commands.CommandError):
"""Generic error with the bot. This can be used to catch all bot errors.""" """Generic error with the bot. This can be used to catch all bot errors."""
pass pass
class ImageResizeTimeoutError(EmoteManagerError, asyncio.TimeoutError): class ImageProcessingTimeoutError(EmoteManagerError, asyncio.TimeoutError):
pass
class ImageResizeTimeoutError(ImageProcessingTimeoutError):
"""Resizing the image took too long.""" """Resizing the image took too long."""
def __init__(self): def __init__(self):
super().__init__('Error: resizing the image took too long.') 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): class HTTPException(EmoteManagerError):
"""The server did not respond with an OK status code.""" """The server did not respond with an OK status code."""
def __init__(self, status): def __init__(self, status):

View File

@ -4,6 +4,7 @@
import asyncio import asyncio
import base64 import base64
import contextlib import contextlib
import functools
import io import io
import logging import logging
import sys import sys
@ -28,42 +29,43 @@ def resize_until_small(image_data: io.BytesIO) -> None:
# so resizing sometimes does more harm than good. # so resizing sometimes does more harm than good.
max_resolution = 128 # pixels max_resolution = 128 # pixels
image_size = size(image_data) 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:
logger.debug('image size too big (%s bytes)', image_size) return
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
try: try:
thumbnail(image_data, (max_resolution, max_resolution)) with wand.image.Image(blob=image_data) as original_image:
except wand.exceptions.CoderError: while True:
raise errors.InvalidImageError 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) with original_image.clone() as resized:
max_resolution //= 2 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: max_resolution //= 2
"""Resize an image in place to no more than max_size pixels, preserving aspect ratio.""" except wand.exceptions.CoderError:
with wand.image.Image(blob=image_data) as image: raise errors.InvalidImageError
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 def convert_to_gif(image_data: io.BytesIO) -> None:
image_data.seek(0) 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): image_data.truncate(0)
"""Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res""" image_data.seek(0)
# https://stackoverflow.com/a/6565988 converted.save(file=image_data)
image_data.seek(0)
old_width, old_height = old_res except wand.exceptions.CoderError:
new_width, new_height = new_res raise errors.InvalidImageError
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
def mime_type_for_image(data): def mime_type_for_image(data):
if data.startswith(b'\x89PNG\r\n\x1a\n'): 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) return fmt.format(mime=mime, data=b64)
def main() -> typing.NoReturn: 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 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()) data = io.BytesIO(sys.stdin.buffer.read())
try: try:
resize_until_small(data) f(data)
except errors.InvalidImageError: except errors.InvalidImageError:
# 2 is used because 1 is already used by python's default error handler # 2 is used because 1 is already used by python's default error handler
sys.exit(2) sys.exit(2)
@ -102,19 +111,19 @@ def main() -> typing.NoReturn:
sys.exit(0) 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( proc = await asyncio.create_subprocess_exec(
sys.executable, '-m', __name__, sys.executable, '-m', __name__, command_name,
stdin=asyncio.subprocess.PIPE, stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE) stderr=asyncio.subprocess.PIPE)
try: 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: except asyncio.TimeoutError:
proc.kill() proc.kill()
raise errors.ImageResizeTimeoutError raise errors.ImageResizeTimeoutError if command_name == 'resize' else errors.ImageConversionTimeoutError
else: else:
if proc.returncode == 2: if proc.returncode == 2:
raise errors.InvalidImageError raise errors.InvalidImageError
@ -123,6 +132,9 @@ async def resize_in_subprocess(image_data: bytes):
return image_data 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): def size(fp):
"""return the size, in bytes, of the data a file-like object represents""" """return the size, in bytes, of the data a file-like object represents"""
with preserve_position(fp): with preserve_position(fp):