Compress images in a separate process

Compression with Pillow can take long, especially with large
clipboard PNG images.
Doing this in a separate process prevents the async event loop from
getting blocked, and allows multiple compression operations to run in
parallel.
This commit is contained in:
miruka 2020-07-21 23:35:16 -04:00
parent 286b7a951a
commit 0088fadddd
3 changed files with 40 additions and 14 deletions

View File

@ -4,7 +4,6 @@
- hflickable: support kinetic scrolling disabler - hflickable: support kinetic scrolling disabler
- settings - settings
- compress png in a thread
- verify upload cancellation - verify upload cancellation
- clipboard preview doesn't update when copied image changes until second time - clipboard preview doesn't update when copied image changes until second time
- Avatar tooltip can get displayed in front of presence menu - Avatar tooltip can get displayed in front of presence menu

View File

@ -597,11 +597,11 @@ class MatrixClient(nio.AsyncClient):
with NamedTemporaryFile(prefix=prefix, suffix=".png") as temp: with NamedTemporaryFile(prefix=prefix, suffix=".png") as temp:
async def get_path() -> Path: async def get_path() -> Path:
with io.BytesIO(image) as inp, io.BytesIO() as buffer: # optimize is too slow for large images
PILImage.open(inp).save(buffer, "PNG") compressed = await utils.compress_image(image, optimize=False)
async with aiofiles.open(temp.name, "wb") as file: async with aiofiles.open(temp.name, "wb") as file:
await file.write(buffer.getvalue()) await file.write(compressed)
return Path(temp.name) return Path(temp.name)
@ -1308,16 +1308,15 @@ class MatrixClient(nio.AsyncClient):
if not small: if not small:
thumb.thumbnail((800, 600)) thumb.thumbnail((800, 600))
with io.BytesIO() as out: if thumb.mode in png_modes:
if thumb.mode in png_modes: thumb_data = await utils.compress_image(thumb)
thumb.save(out, "PNG", optimize=True) mime = "image/png"
mime = "image/png" else:
else: thumb = thumb.convert("RGB")
thumb.convert("RGB").save(out, "JPEG", optimize=True) thumb_data = await utils.compress_image(thumb, "JPEG")
mime = "image/jpeg" mime = "image/jpeg"
thumb_data = out.getvalue() thumb_size = len(thumb_data)
thumb_size = len(thumb_data)
if thumb_size >= len(data) and is_jpg_png and not is_svg: if thumb_size >= len(data) and is_jpg_png and not is_svg:
raise UneededThumbnail() raise UneededThumbnail()

View File

@ -2,6 +2,7 @@
"""Various utilities that are used throughout the package.""" """Various utilities that are used throughout the package."""
import asyncio
import collections import collections
import html import html
import inspect import inspect
@ -9,6 +10,7 @@ import io
import json import json
import sys import sys
import xml.etree.cElementTree as xml_etree # FIXME: bandit warning import xml.etree.cElementTree as xml_etree # FIXME: bandit warning
from concurrent.futures import ProcessPoolExecutor
from datetime import datetime, timedelta from datetime import datetime, timedelta
from enum import Enum from enum import Enum
from enum import auto as autostr from enum import auto as autostr
@ -24,6 +26,8 @@ import aiofiles
import filetype import filetype
from aiofiles.threadpool.binary import AsyncBufferedReader from aiofiles.threadpool.binary import AsyncBufferedReader
from aiofiles.threadpool.text import AsyncTextIOWrapper from aiofiles.threadpool.text import AsyncTextIOWrapper
from PIL import Image as PILImage
from nio.crypto import AsyncDataT as File from nio.crypto import AsyncDataT as File
from nio.crypto import async_generator_from_data from nio.crypto import async_generator_from_data
@ -34,8 +38,11 @@ else:
AsyncOpenFile = Union[AsyncTextIOWrapper, AsyncBufferedReader] AsyncOpenFile = Union[AsyncTextIOWrapper, AsyncBufferedReader]
Size = Tuple[int, int] Size = Tuple[int, int]
BytesOrPIL = Union[bytes, PILImage.Image]
auto = autostr auto = autostr
COMPRESSION_POOL = ProcessPoolExecutor()
class AutoStrEnum(Enum): class AutoStrEnum(Enum):
"""An Enum where auto() assigns the member's name instead of an integer. """An Enum where auto() assigns the member's name instead of an integer.
@ -236,3 +243,24 @@ async def atomic_write(
temp_path.replace(path) temp_path.replace(path)
else: else:
temp_path.unlink() temp_path.unlink()
def _compress(image: BytesOrPIL, fmt: str, optimize: bool) -> bytes:
if isinstance(image, bytes):
pil_image = PILImage.open(io.BytesIO(image))
else:
pil_image = image
with io.BytesIO() as buffer:
pil_image.save(buffer, fmt, optimize=optimize)
return buffer.getvalue()
async def compress_image(
image: BytesOrPIL, fmt: str = "PNG", optimize: bool = True,
) -> bytes:
"""Compress image in a separate process, without blocking event loop."""
return await asyncio.get_event_loop().run_in_executor(
COMPRESSION_POOL, _compress, image, fmt, optimize,
)