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/
|
venv/
|
||||||
|
|
||||||
config.py
|
config.py
|
||||||
|
|
||||||
|
*.png
|
||||||
|
*.gif
|
||||||
|
*.jpg
|
||||||
|
|
|
@ -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
|
||||||
|
async with context.typing():
|
||||||
archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES)
|
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
|
||||||
|
|
|
@ -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):
|
||||||
|
|
|
@ -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:
|
||||||
|
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('image size too big (%s bytes)', image_size)
|
||||||
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
|
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
|
||||||
|
|
||||||
try:
|
with original_image.clone() as resized:
|
||||||
thumbnail(image_data, (max_resolution, max_resolution))
|
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:
|
except wand.exceptions.CoderError:
|
||||||
raise errors.InvalidImageError
|
raise errors.InvalidImageError
|
||||||
|
|
||||||
image_size = size(image_data)
|
def convert_to_gif(image_data: io.BytesIO) -> None:
|
||||||
max_resolution //= 2
|
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.truncate(0)
|
||||||
image_data.seek(0)
|
image_data.seek(0)
|
||||||
image.save(file=image_data)
|
converted.save(file=image_data)
|
||||||
|
|
||||||
# allow resizing the original image more than once for memory profiling
|
|
||||||
image_data.seek(0)
|
image_data.seek(0)
|
||||||
|
except wand.exceptions.CoderError:
|
||||||
def scale_resolution(old_res, new_res):
|
raise errors.InvalidImageError
|
||||||
"""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
|
|
||||||
|
|
||||||
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):
|
||||||
|
|
Loading…
Reference in a new issue