commands.py

"""Command implementations."""

import abc
import asyncio
import asyncio.subprocess
import datetime
import logging
import pathlib
import platform
import re
import shlex
import shutil
import subprocess
import tempfile

import time
import typing
from pathlib import Path

import attr
import psutil

from . import worker

command_handlers = {}  # type: typing.Mapping[str, typing.Type['AbstractCommand']]

# Some type declarations.
Settings = typing.MutableMapping[str, typing.Any]
# This is the type of the 2nd arg for instanceof(a, b)
InstanceOfType = typing.Union[type, typing.Tuple[typing.Union[type, typing.Tuple[typing.Any, ...]],
                                                 ...]]

# Timeout of subprocess.stdout.readline() call.
SUBPROC_READLINE_TIMEOUT = 3600  # seconds

MERGE_EXR_PYTHON = """\
import bpy
scene = bpy.context.scene
nodes = scene.node_tree.nodes
image1 = bpy.data.images.load("%(input1)s")
image2 = bpy.data.images.load("%(input2)s")

nodes["image1"].image = image1
nodes["image2"].image = image2
nodes["weight1"].outputs[0].default_value = %(weight1)i
nodes["weight2"].outputs[0].default_value = %(weight2)i

nodes["output"].base_path = "%(tmpdir)s"
scene.render.resolution_x, scene.render.resolution_y = image1.size
scene.render.tile_x, scene.render.tile_y = image1.size
scene.render.filepath = "%(tmpdir)s/preview.jpg"

bpy.ops.render.render(write_still=True)
"""

log = logging.getLogger(__name__)


def command_executor(cmdname):
    """Class decorator, registers a command executor."""

    def decorator(cls):
        assert cmdname not in command_handlers

        command_handlers[cmdname] = cls
        cls.command_name = cmdname
        return cls

    return decorator


class CommandExecutionError(Exception):
    """Raised when there was an error executing a command."""
    pass


@attr.s
class AbstractCommand(metaclass=abc.ABCMeta):
    """Command executor.

    This class (or any of its subclasses) should not directly set the task status.
    This should be left to the Worker.
    """
    worker = attr.ib(validator=attr.validators.instance_of(worker.FlamencoWorker),
                     repr=False)
    task_id = attr.ib(validator=attr.validators.instance_of(str))
    command_idx = attr.ib(validator=attr.validators.instance_of(int))

    # Set by @command_executor
    command_name = ''

    # Set by __attr_post_init__()
    identifier = attr.ib(default=None, init=False,
                         validator=attr.validators.optional(attr.validators.instance_of(str)))
    _log = attr.ib(init=False, default=logging.getLogger('AbstractCommand'),
                   validator=attr.validators.instance_of(logging.Logger))

    def __attrs_post_init__(self):
        self.identifier = '%s.(task_id=%s, command_idx=%s)' % (
            self.command_name,
            self.task_id,
            self.command_idx)
        self._log = log.getChild(self.identifier)

    async def run(self, settings: Settings) -> bool:
        """Runs the command, parsing output and sending it back to the worker.

        Returns True when the command was succesful, and False otherwise.
        """

        verr = self.validate(settings)
        if verr is not None:
            self._log.warning('Error in settings: %s', verr)
            await self.worker.register_log('%s: Error in settings: %s', self.identifier, verr)
            await self.worker.register_task_update(
                task_status='failed',
                activity='%s: Invalid settings: %s' % (self.identifier, verr),
            )
            return False

        await self.worker.register_log('%s: Starting' % self.command_name)
        await self.worker.register_task_update(
            activity='starting %s' % self.command_name,
            current_command_idx=self.command_idx,
            command_progress_percentage=0
        )

        try:
            await self.execute(settings)
        except CommandExecutionError as ex:
            # This is something we threw ourselves, and there is no need to log the traceback.
            self._log.warning('Error executing: %s', ex)
            await self._register_exception(ex)
            return False
        except asyncio.CancelledError as ex:
            self._log.warning('Command execution was canceled')
            raise
        except Exception as ex:
            # This is something unexpected, so do log the traceback.
            self._log.exception('Error executing.')
            await self._register_exception(ex)
            return False

        await self.worker.register_log('%s: Finished' % self.command_name)
        await self.worker.register_task_update(
            activity='finished %s' % self.command_name,
            current_command_idx=self.command_idx,
            command_progress_percentage=100
        )

        return True

    async def abort(self):
        """Aborts the command. This may or may not be actually possible.

        A subprocess that's started by this command will be killed.
        However, any asyncio coroutines that are not managed by this command
        (such as the 'run' function) should be cancelled by the caller.
        """

    async def _register_exception(self, ex: Exception):
        """Registers an exception with the worker, and set the task status to 'failed'."""

        await self.worker.register_log('%s: Error executing: %s' % (self.identifier, ex))
        await self.worker.register_task_update(
            activity='%s: Error executing: %s' % (self.identifier, ex),
        )

    @abc.abstractmethod
    async def execute(self, settings: Settings) -> None:
        """Executes the command.

        An error should be indicated by an exception.
        """

    def validate(self, settings: Settings) -> typing.Optional[str]:
        """Validates the settings for this command.

        If there is an error, a description of the error is returned.
        If the settings are valid, None is returned.

        By default all settings are considered valid.
        """

        return None

    def _setting(self, settings: Settings, key: str, is_required: bool,
                 valtype: InstanceOfType = str,
                 default: typing.Any = None) \
            -> typing.Tuple[typing.Any, typing.Optional[str]]:
        """Parses a setting, returns either (value, None) or (None, errormsg)"""

        try:
            value = settings[key]
        except KeyError:
            if is_required:
                return None, 'Missing "%s"' % key
            settings.setdefault(key, default)
            return default, None