play-daemon-threaded/pipeline/handlers/poster.py

from os import path, remove, rename
from shutil import copy2

import ffmpeg

from .handler import Handler
from .transcode import TranscodeHandler
from ..exceptions import ValidationException


def make_poster(inpath, outdir):
    """
    Create a poster from the video file at inpath.
    """
    indir, infile = path.split(inpath)
    infile_noext, _ = path.splitext(infile)
    outfile = f"{infile_noext}-poster.jpg"
    outpath = path.join(outdir, outfile)

    ffmpeg.input(inpath) \
          .output(outpath, vframes=1) \
          .run(quiet=True)
    return (outdir, outfile)


def _needs_poster(source):
    """
    Returns True if a poster should be generated for this source,
    otherwise False.

    Poster should be generated if:
    - a poster key has been passed with falsy value
    - a video has been uploaded without a corresponding poster file
    """
    if 'poster' in source and not source['poster']:
        return True
    if 'video' in source and not source.get('poster', None):
        return True
    return False


@Handler.register
class PosterHandler(Handler):
    """
    This class generates poster files and imports uploaded ones.
    """
    @classmethod
    def wants(cls, jobspec, existing_package):
        """
        Return True if this handler wants to process this jobspec.
        Raises an exception if the job is wanted but doesn't pass validation.

        A job is wanted if any of:
         - a 'poster' key exists in a source
         - a video has been uploaded for a source without a
           corresponding poster
        """
        if 'sources' in jobspec:
            # A poster has been uploaded
            uploaded = {name: source['poster']
                        for name, source in jobspec['sources'].items()
                        if 'poster' in source and source['poster']}
            # A poster is to be generated
            generate = [_needs_poster(source)
                        for source in jobspec['sources'].values()]
            if any(uploaded) or any(generate):
                return cls._validate(jobspec, existing_package, uploaded)
        return False

    @classmethod
    def _validate(cls, jobspec, existing_package, uploaded):
        """
        Return True if the job is valid for this handler.

        Validity requirements are:
         - 'upload_dir' must exist if a 'poster' value is truthy
         - all truthy 'poster' values must point to valid files
        """
        super()._validate(jobspec, existing_package)
        if uploaded:
            if 'upload_dir' not in jobspec:
                raise ValidationException("upload_dir missing")
        for name, poster in uploaded.items():
            if not path.isfile(path.join(jobspec['upload_dir'],
                                         poster)):
                raise ValidationException(
                    f"poster file invalid for source {name}")
        return True

    @classmethod
    def apply_after(cls, handlerclass):
        """
        Return True if this handler must be applied after
        the supplied handler.

        This handler must be applied after the transcode handler,
        so posters can be generated from the right files.
        """
        if issubclass(handlerclass, TranscodeHandler):
            return True
        return False

    def _handle(self, jobspec, existing_package, tempdir):
        """
        Return a function to apply changes to the stored package.

        The returned function generates posters as appropriate, moves poster
        files into the package's basedir and updates the package metadata.
        Replaced posters are deleted.
        """
        def apply_func(package):
            uploaded = {name: source['poster']
                        for name, source in jobspec['sources'].items()
                        if 'poster' in source and source['poster']}
            for name, poster in uploaded.items():
                if name in package['sources']:
                    oldposter = package['sources'][name].get('poster', None)
                    if oldposter:
                        remove(path.join(package.basedir, oldposter))
                else:
                    # Initialize a new source
                    # set playAudio to False by default so we don't conflict
                    # with the audio handler
                    package['sources'][name] = {'playAudio': False}
                # Don't rename - the file may be used by multiple sources.
                posterfile = path.basename(poster)
                copy2(path.join(jobspec['upload_dir'], poster),
                      path.join(package.basedir, posterfile))
                package['sources'][name]['poster'] = posterfile

            generate = [name for name, source in jobspec['sources'].items()
                        if _needs_poster(source)]
            for name in generate:
                if name not in package['sources']:
                    # A job could request thumbnail generation for a stream
                    # that doesn't exist. Ignore it, since there is nothing
                    # to generate from.
                    continue
                packagesource = package['sources'][name]

                oldposter = packagesource.get('poster', None)
                if oldposter:
                    remove(path.join(package.basedir, oldposter))

                # We generate the poster from the largest resolution available
                largest = str(max(int(height) for height
                                  in packagesource['video'].keys()))
                videoname = packagesource['video'][largest]
                _, postername = make_poster(path.join(package.basedir,
                                                      videoname),
                                            tempdir)
                # We can rename here because the files are unique.
                rename(path.join(tempdir, postername),
                       path.join(package.basedir, postername))
                packagesource['poster'] = postername

        return apply_func