photolib/photoapp/thumb.py

import os
import sys
import traceback
from time import time
from collections import defaultdict
from multiprocessing import Process
from PIL import Image, ImageOps
import tempfile
from shutil import copyfileobj
from contextlib import closing
import logging


class ThumbGenerator(object):
    def __init__(self, library, storage):
        self.library = library
        self.storage = storage
        self._failed_thumbs_cache = defaultdict(dict)

    def get_datedir_path(self, date):
        """
        Return a path like 2018/3/31 given a datetime object representing the same date
        """
        return os.path.join(str(date.year), str(date.month), str(date.day))

    def make_thumb(self, photo, style):
        """
        Create a thumbnail of the given photo, scaled/cropped to the given named style
        :return: local path to thumbnail file or None if creation failed or was blocked
        """
        # style tuples: max x, max y, rotate ok)
        # rotate ok means x and y maxes can be swapped if it fits the image's aspect ratio better
        styles = {"tiny": (80, 80, False),
                  "small": (100, 100, False),
                  "feed": (250, 250, False),
                  "preview": (1024, 768, True),
                  "big": (2048, 1536, True)}
        dest = os.path.join("thumbs", style, photo.uuid[0:2], "{}.jpg".format(photo.uuid))

        if self.storage.exists(dest):
            return self.storage.open(dest, "rb")
        if photo.width is None:  # todo better detection of images that PIL can't open
            return None
        if photo.uuid in self._failed_thumbs_cache[style]:
            return None

        thumb_width, thumb_height, flip_ok = styles[style]
        i_width = photo.width
        i_height = photo.height
        im_is_rotated = photo.orientation % 2 != 0 or i_height > i_width

        if im_is_rotated and flip_ok:
            thumb_width, thumb_height = thumb_height, thumb_width

        thumb_width = min(thumb_width, i_width if i_width > 0 else 999999999)  # TODO do we even have photo.width if PIL can't read the image?
        thumb_height = min(thumb_height, i_height if i_height > 0 else 999999999)  # TODO this seems bad

        # TODO have the subprocess download the file
        # TODO thundering herd
        with tempfile.TemporaryDirectory() as tmpdir:
            fthumblocal = tempfile.NamedTemporaryFile(delete=True)
            fpath = os.path.join(tmpdir, "image")
            with self.library.storage.open(photo.path, 'rb') as fsrc:
                with open(fpath, 'wb') as ftmpdest:
                    copyfileobj(fsrc, ftmpdest)

            p = Process(target=self.gen_thumb, args=(fpath, fthumblocal.name, thumb_width, thumb_height, photo.orientation))
            p.start()
            p.join()
            if p.exitcode != 0:
                self._failed_thumbs_cache[style][photo.uuid] = True  # dont retry failed generations
                return None

            with closing(self.storage.open(dest, 'wb')) as fdest:
                copyfileobj(fthumblocal, fdest)

        fthumblocal.seek(0)
        return fthumblocal

    @staticmethod
    def gen_thumb(src_img, dest_img, width, height, rotation):
        try:
            start = time()
            # TODO lock around the dir creation
            os.makedirs(os.path.split(dest_img)[0], exist_ok=True)
            image = Image.open(src_img)
            image = image.rotate(90 * rotation, expand=True)
            thumb = ImageOps.fit(image, (width, height), Image.ANTIALIAS)
            thumb.save(dest_img, 'JPEG')
            logging.info("Generated {} in {}s".format(dest_img, round(time() - start, 4)))
        except Exception:
            traceback.print_exc()
            if os.path.exists(dest_img):
                os.unlink(dest_img)
            sys.exit(1)