batch_upload_thumbnails_baseaction.py

# :coding: utf-8
import logging
import os

import ftrack_api

from ftrack_action_handler.action import BaseAction


class ImportThumbnailsBaseAction(BaseAction):
    """Batch import thumbnails to a project."""

    label = "Batch Import Thumbnails"
    identifier = "ftrack.batch.import.thumbnails.base"
    description = "Batch import thumbnails from folder to project."

    ui_label = (
        "The action will batch import thumbnails to the selected project.\n\n"
        "Specify a *folder path* to a folder containing the images.\n\n"
        "The images should be named to match the entity path in ftrack.\n\n"
        "For example:\n\n"
        "    0010.png\n"
        "    0010.010.png\n"
        "    0010.010.generic.png\n"
        "\n"
        "This will set the thumbnail for the *sequence*, *shot*"
        " and the *generic task*."
    )


    def discover(self, session, entities, event):
        if not self.validate_selection(entities):
            return super(ImportThumbnailsBaseAction, self).discover(
                session, entities,  event
            )

        return True

    def interface(self, session, entities, event):
        input_values = event["data"].get("values", {})
        #TODO: display error message when invalid input path
        if not self.validate_input(input_values):
            return [
                {
                    "type": "label",
                    "value": self.ui_label
                },
                {
                    "type": "text",
                    "label": "Folder path",
                    "name": "folder_path"
                }
            ]

    def launch(self, session, entities, event):
        if not self.validate_selection(entities):
            self.logger.warning("Selection is not valid => Aborting action")
            return

        input_values = event["data"]["values"]

        # Project entity
        project_id = entities[0][1] # entity ID
        project_entity = self.session.query(
            "Project where id is {}".format(project_id)
        ).one()
        project_name = project_entity["full_name"]
        self.logger.info("Project name: {}".format(project_name))

        # Files
        folder_path = input_values["folder_path"]
        files = self.get_files(folder_path, project_name)

        # Batch process
        try:
            self.upload_files_as_thumbnails(project_entity, files)
            self.session.commit()
        except Exception as exc:
            #TODO: filter exception type?
            self.logger.error("Error during process: {}".format(str(exc)))
            self.session.rollback()
            raise

        message = "Batch Import Thumbnails action completed successfully"
        self.logger.info(message)
        return {
            "success": True,
            "message": message
        }

    def upload_files_as_thumbnails(self, project, files):
        """Upload the files as thumbnails."""

        # Get project entity and all its children entities
        entities = {
            project["full_name"]: project
        }
        entities.update(
            self.get_child_elements_from_entity(self.session, project)
        )

        # Find matching entities
        filtered_entities = self.filter_entities(entities, files)

        self.session._configure_locations()
        server_location = self.session.query(
            "Location where name is \"ftrack.server\""
        ).one()

        for ftrack_path in filtered_entities:
            file = files[ftrack_path]
            entities = filtered_entities[ftrack_path]
            for entity in entities:
                self.logger.debug(
                    "Setting thumbnail '{}' for entity '{}'".format(file,
                                                                    entity)
                )
                thumbnail_component = self.session.create_component(
                    file,
                    dict(name="thumbnail"),
                    location=server_location
                )
                entity["thumbnail"] = thumbnail_component

    @staticmethod
    def validate_selection(entities):
        """Return True if *entities* is valid.

        Selection is valid if consists of a single project entity.
        """

        if (
            not entities or
            len(entities) != 1 or
            entities[0][0].lower() != "project"
        ):
            return False

        return True

    @staticmethod
    def validate_input(input_values):
        """Validate the input parameters."""

        if not input_values:
            return False

        folder_path = input_values.get("folder_path", None)
        if not folder_path:
            return False

        if not os.path.isdir(folder_path):
            return False

        return True

    #TODO: add option to have full path (=> add project_name)
    # (also need to do it to entity paths)
    @staticmethod
    def get_files(folder_path, project_name):
        """Return a list of tuples with ftrack path and filepaths."""
        all_files = {}
        for filename in os.listdir(folder_path):
            absolute_file_path = os.path.join(folder_path, filename)
            if os.path.isfile(absolute_file_path):
                ftrack_path, _ = os.path.splitext(filename)
                all_files[ftrack_path] = absolute_file_path

        return all_files

    @staticmethod
    def filter_entities(entities, files):
        filtered_entities = {}
        for entity_path in entities:
            entity_path_lower = entity_path.lower()
            for ftrack_path in files:
                if entity_path_lower.endswith(ftrack_path.lower()):
                    # Stop at first match
                    break
            else:
                continue

            # Found a match
            if ftrack_path not in filtered_entities:
                filtered_entities[ftrack_path] = []
            filtered_entities[ftrack_path].append(entities[entity_path])

        return filtered_entities

    @classmethod
    def get_child_elements_from_entity(cls, session, entity):
        children = {}
        for child in entity["children"]:
            ancestors = []
            for ancestor in child["ancestors"]:
                ancestors.append(ancestor["name"])
            entity_path = ".".join(ancestors + [child["name"]])

            children[entity_path] = child
            if entity.entity_type.lower() != "task":
                children.update(
                    cls.get_child_elements_from_entity(
                        session, child
                    )
                )

        return children


def register(session, **kw):
    if not isinstance(session, ftrack_api.session.Session):
        return

    action = ImportThumbnailsBaseAction(session)
    action.register()


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    session = ftrack_api.Session()
    register(session)

    session.event_hub.connect()

    logging.info("Registered actions and listening for events."
                 " Use Ctrl-C to abort.")
    try:
        session.event_hub.wait()
    except KeyboardInterrupt:
        logging.info("Process terminated by Ctrl-C.")