view sat/tools/image.py @ 3252:54934ee3f69c

tools (image): added a guess_type method to guess media type: media type is guessed first from filename, then if it's not enough, the file is opened to try to guess type. If it still can be determined, None is returned.
author Goffi <goffi@goffi.org>
date Tue, 14 Apr 2020 20:29:37 +0200
parents 4fbea7f1e012
children f300d78f08f3
line wrap: on
line source

#!/usr/bin/env python3

# SàT: an XMPP client
# Copyright (C) 2009-2020 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Methods to manipulate images"""

import tempfile
import mimetypes
from PIL import Image
from pathlib import Path
from twisted.internet import threads


def check(host, path, max_size=None):
    """Analyze image and return a report

    report will indicate if image is too large, and the recommended new size if this is
    the case
    @param host: SàT instance
    @param path(str, pathlib.Path): image to open
    @param max_size(tuple[int, int]): maximum accepted size of image
        None to use value set in config
    @return dict: report on image, with following keys:
        - too_large: true if image is oversized
        - recommended_size: if too_large is True, recommended size to use
    """
    report = {}
    image = Image.open(path)
    if max_size is None:
        max_size = tuple(host.memory.getConfig(None, "image_max", (1200, 720)))
    if image.size > max_size:
        report['too_large'] = True
        if image.size[0] > max_size[0]:
            factor = max_size[0] / image.size[0]
            if image.size[1] * factor > max_size[1]:
                factor = max_size[1] / image.size[1]
        else:
            factor = max_size[1] / image.size[1]
        report['recommended_size'] = [int(image.width*factor), int(image.height*factor)]
    else:
        report['too_large'] = False

    return report


def _resizeBlocking(image_path, new_size, dest=None):
    im_path = Path(image_path)
    im = Image.open(im_path)
    resized = im.resize(new_size, Image.LANCZOS)

    if dest is None:
        dest = tempfile.NamedTemporaryFile(suffix=im_path.suffix, delete=False)
    elif isinstance(dest, Path):
        dest = dest.open('wb')

    with dest as f:
        resized.save(f, format=im.format)

    return Path(f.name)


def resize(image_path, new_size, dest=None):
    """Resize an image to a new file, and return its path

    @param image_path(str, Path): path of the original image
    @param new_size(tuple[int, int]): size to use for new image
    @param dest(None, Path, file): where the resized image must be stored, can be:
        - None: use a temporary file
        - Path: path to the file to create/overwrite
        - file: a file object which must be opened for writing in binary mode
    @return (Path): path of the resized file.
        The image at this path should be deleted after use
    """
    return threads.deferToThread(_resizeBlocking, image_path, new_size, dest)


def guess_type(source):
    """Guess image media type

    @param source(str, Path, file): image to guess type
    @return (str, None): media type, or None if we can't guess
    """
    if isinstance(source, str):
        source = Path(source)

    if isinstance(source, Path):
        # we first try to guess from file name
        media_type = mimetypes.guess_type(source, strict=False)[0]
        if media_type is not None:
            return media_type

    # file name is not enough, we try to open it
    img = Image.open(source)
    try:
        return Image.MIME[img.format]
    except KeyError:
        return None