o
/
EmoteManager
mirror of https://github.com/uhIgnacio/EmoteManager.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

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

52
cogs/emote.py
View File

@ -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

9
utils/errors.py
View File

@ -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):

86
utils/image.py
View File

@ -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):