mirror of
https://github.com/uhIgnacio/EmoteManager.git
synced 2024-08-15 02:23:13 +00:00
move image resizing to a subprocess
This commit is contained in:
parent
943ca5d8a6
commit
aae88068fc
2 changed files with 121 additions and 43 deletions
|
@ -1,9 +1,8 @@
|
|||
#!/usr/bin/env python3
|
||||
# encoding: utf-8
|
||||
|
||||
import io
|
||||
import imghdr
|
||||
import asyncio
|
||||
import cgi
|
||||
import logging
|
||||
import weakref
|
||||
import traceback
|
||||
|
@ -186,29 +185,31 @@ class Emotes(commands.Cog):
|
|||
return emote
|
||||
|
||||
async def fetch_emote(self, url):
|
||||
# credits to @Liara#0001 (ID 136900814408122368) for most of this part
|
||||
# https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L217-228
|
||||
async with self.http.head(url, timeout=5) as response:
|
||||
def validate_headers(response):
|
||||
if response.reason != 'OK':
|
||||
raise errors.HTTPException(response.status)
|
||||
if response.headers.get('Content-Type') not in ('image/png', 'image/jpeg', 'image/gif'):
|
||||
# some dumb servers also send '; charset=UTF-8' which we should ignore
|
||||
mimetype, options = cgi.parse_header(response.headers.get('Content-Type', ''))
|
||||
if mimetype not in {'image/png', 'image/jpeg', 'image/gif'}:
|
||||
raise errors.InvalidImageError
|
||||
|
||||
async with self.http.get(url) as response:
|
||||
if response.reason != 'OK':
|
||||
raise errors.HTTPException(response.status)
|
||||
return io.BytesIO(await response.read())
|
||||
try:
|
||||
async with self.http.head(url, timeout=5) as response:
|
||||
validate_headers(response)
|
||||
except aiohttp.ServerDisconnectedError as exception:
|
||||
validate_headers(exception.message)
|
||||
|
||||
async def create_emote_from_bytes(self, guild, name, author_id, image_data: io.BytesIO, *, reason=None):
|
||||
# resize_until_small is normally blocking, because wand is.
|
||||
# run_in_executor is magic that makes it non blocking somehow.
|
||||
# also, None as the executor arg means "use the loop's default executor"
|
||||
image_data = await self.bot.loop.run_in_executor(None, utils.image.resize_until_small, image_data)
|
||||
async with self.http.get(url) as response:
|
||||
validate_headers(response)
|
||||
return await response.read()
|
||||
|
||||
async def create_emote_from_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None):
|
||||
image_data = await utils.image.resize_in_subprocess(image_data)
|
||||
if reason is None:
|
||||
reason = f'Created by {utils.format_user(self.bot, author_id)}'
|
||||
return await guild.create_custom_emoji(
|
||||
name=name,
|
||||
image=image_data.read(),
|
||||
image=image_data,
|
||||
reason=reason)
|
||||
|
||||
@commands.command(aliases=('delete', 'delet', 'rm'))
|
||||
|
|
131
utils/image.py
Normal file → Executable file
131
utils/image.py
Normal file → Executable file
|
@ -1,18 +1,27 @@
|
|||
#!/usr/bin/env python3
|
||||
# encoding: utf-8
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import io
|
||||
import logging
|
||||
import contextlib
|
||||
import sys
|
||||
import typing
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from wand.image import Image
|
||||
except ImportError:
|
||||
import wand.image
|
||||
except (ImportError, OSError):
|
||||
logger.warn('Failed to import wand.image. Image manipulation functions will be unavailable.')
|
||||
else:
|
||||
import wand.exceptions
|
||||
|
||||
from utils import errors
|
||||
|
||||
def resize_until_small(image_data: io.BytesIO):
|
||||
"""If the image_data is bigger than 256KB, resize it until it's not"""
|
||||
def resize_until_small(image_data: io.BytesIO) -> None:
|
||||
"""If the image_data is bigger than 256KB, resize it until it's not."""
|
||||
# It's important that we only attempt to resize the image when we have to,
|
||||
# ie when it exceeds the Discord limit of 256KiB.
|
||||
# Apparently some <256KiB images become larger when we attempt to resize them,
|
||||
|
@ -21,43 +30,108 @@ def resize_until_small(image_data: io.BytesIO):
|
|||
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 %s*%s pixels', max_resolution, max_resolution)
|
||||
image_data = thumbnail(image_data, (max_resolution, max_resolution))
|
||||
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
|
||||
|
||||
try:
|
||||
thumbnail(image_data, (max_resolution, max_resolution))
|
||||
except wand.exceptions.CoderError:
|
||||
raise errors.InvalidImageError
|
||||
|
||||
image_size = size(image_data)
|
||||
max_resolution //= 2
|
||||
return image_data
|
||||
|
||||
def size(data: io.BytesIO):
|
||||
"""return the size, in bytes, of the data a BytesIO object represents"""
|
||||
with preserve_position(data):
|
||||
data.seek(0, io.SEEK_END)
|
||||
return data.tell()
|
||||
|
||||
def thumbnail(image_data: io.BytesIO, max_size=(128, 128)):
|
||||
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."""
|
||||
# Credit to @Liara#0001 (ID 136900814408122368)
|
||||
# https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L243-247
|
||||
image = Image(blob=image_data)
|
||||
image.resize(*scale_resolution((image.width, image.height), max_size))
|
||||
# we create a new buffer here because there's wand errors otherwise.
|
||||
# specific error:
|
||||
# MissingDelegateError: no decode delegate for this image format `' @ error/blob.c/BlobToImage/353
|
||||
out = io.BytesIO()
|
||||
image.save(file=out)
|
||||
out.seek(0)
|
||||
return out
|
||||
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
|
||||
image_data.seek(0)
|
||||
|
||||
def scale_resolution(old_res, new_res):
|
||||
# https://stackoverflow.com/a/6565988
|
||||
"""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):
|
||||
if data.startswith(b'\x89PNG\r\n\x1a\n'):
|
||||
return 'image/png'
|
||||
elif data.startswith(b'\xFF\xD8') and data.rstrip(b'\0').endswith(b'\xFF\xD9'):
|
||||
return 'image/jpeg'
|
||||
elif data.startswith((b'GIF87a', b'GIF89a')):
|
||||
return 'image/gif'
|
||||
elif data.startswith(b'RIFF') and data[8:12] == b'WEBP':
|
||||
return 'image/webp'
|
||||
else:
|
||||
raise errors.InvalidImageError
|
||||
|
||||
def image_to_base64_url(data):
|
||||
fmt = 'data:{mime};base64,{data}'
|
||||
mime = mime_type_for_image(data)
|
||||
b64 = base64.b64encode(data).decode('ascii')
|
||||
return fmt.format(mime=mime, data=b64)
|
||||
|
||||
def main() -> typing.NoReturn:
|
||||
"""resize an image from stdin and write the resized version to stdout."""
|
||||
import sys
|
||||
|
||||
data = io.BytesIO(sys.stdin.buffer.read())
|
||||
try:
|
||||
resize_until_small(data)
|
||||
except errors.InvalidImageError:
|
||||
# 2 is used because 1 is already used by python's default error handler
|
||||
sys.exit(2)
|
||||
|
||||
stdout_write = sys.stdout.buffer.write # getattr optimization
|
||||
|
||||
while True:
|
||||
buf = data.read(16 * 1024)
|
||||
if not buf:
|
||||
break
|
||||
|
||||
stdout_write(buf)
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
async def resize_in_subprocess(image_data: bytes):
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
sys.executable, '-m', __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)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
raise errors.ImageResizeTimeoutError
|
||||
else:
|
||||
if proc.returncode == 2:
|
||||
raise errors.InvalidImageError
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(err.decode('utf-8') + f'Return code: {proc.returncode}')
|
||||
|
||||
return image_data
|
||||
|
||||
def size(fp):
|
||||
"""return the size, in bytes, of the data a file-like object represents"""
|
||||
with preserve_position(fp):
|
||||
fp.seek(0, io.SEEK_END)
|
||||
return fp.tell()
|
||||
|
||||
class preserve_position(contextlib.AbstractContextManager):
|
||||
def __init__(self, fp):
|
||||
self.fp = fp
|
||||
|
@ -65,3 +139,6 @@ class preserve_position(contextlib.AbstractContextManager):
|
|||
|
||||
def __exit__(self, *excinfo):
|
||||
self.fp.seek(self.old_pos)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
Loading…
Reference in a new issue