download_all_components_action

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

"""
Module download_all_components_from_entity: This module registers an action
that allows the user to download all components from an entity in ftrack.
"""

import argparse
import getpass
import logging
import sys
import traceback
import urllib2
import urlparse
import ftrack
import threading
import os


# add accessor to path and import it
accessorPath = os.path.join(os.path.dirname(
    os.path.dirname(
        os.path.dirname(os.path.abspath(__file__)))), 'resolver')

sys.path.append(accessorPath)

# noinspection PyUnresolvedReferences
import frostburnS3Accessor as fbS3


class DownloadAllComponentsFromEntityAction(ftrack.Action):
    """
    This ``ftrack.Action`` subclass will allow the user to download all
    components from an entity.
    """

    # Action identifier
    identifier = 'download-all-components-from-entity'

    # Action label
    label = 'Download All Components'

    # Current user from OS environment
    currentUser = getpass.getuser()


    def discover(self, event):
        """
        This method is invoked when the user brings up the Actions menu from the
        web UI. This action will only return the config when triggered on a single
        ``ftrack.AssetVersion`` object.
        """

        data = event['data']

        selection = data.get('selection', [])
        self.logger.info('Got selection: {0}'.format(selection))

        # If selection contains more than one item return early
        if len(selection) != 1 or selection[0]['entityType'] != 'assetversion':
            return

        # Return data to show in web UI
        return {
            'items' : [{
                'label' : self.label,
                'actionIdentifier' : self.identifier
            }]
        }


    def launch( self, event ):
        """
        This method is invoked when the user executes the Action. This brings up
        a UI to choose where to save the components to.
        """

        # This returns a dictionary containing the data returned from the server.
        data = event['data']

        # Get list of all component urls attached to this AssetVersion
        components = self.getAllComponentUrls(data)

        # Check for no Components attached to AssetVersion
        if not components:
            self.logger.error('### No components were found on the AssetVersion!')

            # Return early with status message
            return {
                'success' : False,
                'message' : 'No Components found on this Asset Version!'
            }

        # Format dict to be used in the Custom Actions UI
        uiItems = {
            'items' : [
                {
                    'value' : '## Download Individual Components ##',
                    'type' : 'label'
                }
            ]
        }

        # For each url, add a new download link label
        for name, url in components.items():
            urlLinkItem = {
                'value' : 'Download: [{0}]({1})'.format(name, url),
                'type' : 'label'
            }

            uiItems['items'].append(urlLinkItem)

        # Add field for user-specified location to save all Components to

        saveToDirectoryHeader = {
            'value' : '## Save All Components To Directory ##',
            'type' : 'label'
        }

        uriSaveDirectoryItem = {
            'label' : 'Save all components to',
            'type' : 'text',
            'value' : 'Enter path to directory...',
            'name' : 'saveDirectory'
        }

        uiItems['items'].append(saveToDirectoryHeader)
        uiItems['items'].append(uriSaveDirectoryItem)

        # Check to see if values have been returned
        if 'values' in data:
            self.logger.info('Action executed')

            # Open file browser to choose save directory
            # Hide root widget
            # rootDialog = Tkinter.Tk()
            # rootDialog.withdraw()

            # Get user-specified location
            # saveDirectory = tkFileDialog.askdirectory()

            saveDirectory = data['values']['saveDirectory']

            # check for valid directory
            if not os.path.isdir(saveDirectory):
                self.logger.error('### The directory chosen is invalid: {0}'
                                  .format(saveDirectory))

                # return early and publish failure message
                return {
                    'success' : False,
                    'message' : 'Directory chosen is invalid!'
                }

            elif not os.access(saveDirectory, os.W_OK):
                self.logger.error('### You do not have write access to: {0}'
                                  .format(saveDirectory))

                # return early and publish failure message
                return {
                    'success' : False,
                    'message' : 'No permission to write to directory!'
                }

            self.logger.info('identifier')

            # Save all Components to the specified directory
            initializeDownloads(urls=components.values(),
                                fileNames=None,
                                saveDirectory=saveDirectory)

            # Publish success message to UI
            return {
                'success' : True,
                'message' : 'Successfully executed action: {0}!'.format(self.label)
            }

        # return dict with formatted data for displaying the DAC UI
        return uiItems


    def register(self):
        """
        This overloads the base register() function to only subscribe to receive
        certain events that have user information in them.

        Used to filter on users that initiated the action.
        """

        self.logger.debug('register() called from: {0}'.format(__name__))

        # only subscribe/launch if this action's username match.
        ftrack.EVENT_HUB.subscribe(
            subscription='topic=ftrack.action.discover and source.user.username={0}'
                .format(self.currentUser),
            callback=self.discover
        )

        ftrack.EVENT_HUB.subscribe(
            subscription='topic=ftrack.action.launch and data.actionIdentifier={0} '
                         'and source.user.username={1}'
                .format(self.identifier, self.currentUser),
            callback=self.launch
        )


    def getAllComponentUrls(self, data):
        """
        This method retrieves all available urls to ``ftrack.Component`` objects
        from a ``ftrack.AssetVersion``. Returns the results as a ``list``.
        :param data: ``dict`` containing the data returned from the ftrack server.
        :type data: ``dict``
        :return: ``namedtuple`` containing ``list`` objects of Component names and urls
        :rtype: ``namedtuple``
        """

        # make dictionary to store all data
        componentData = {}

        # Query AssetVersion selection
        selection = data.get('selection', [])

        # check for non-existent assetVersion selection
        if not selection:
            return

        for entity in selection:
            # get assetVersion instance from id
            assetVersion = ftrack.AssetVersion(id=entity['entityId'])

            # get all components attached to assetVersion
            components = assetVersion.getComponents()

            for component in components:

                assert isinstance(component, ftrack.Component), \
                    '### {0} is not a ftrack Component!'.format(component)

                # first attempt to get local component path
                componentPath = component.getFilesystemPath()

                # get Component name to store in namedtuple
                componentName = component.getName()

                if not componentPath:
                    self.logger.info('Component {0} location is not local, retrieving from S3...'
                        .format(component.getName()))

                    # get resource identifier for returning path from accessor
                    componentId = component.getId()
                    componentResourceIdentifier = component.getResourceIdentifier()

                    # get location instance from the assetVersion id
                    location = ftrack.pickLocation(componentId)
                    accessor = location.getAccessor()

                    if isinstance(accessor, ftrack.S3Accessor):
                        # re-class to extended S3Accessor
                        fbS3.FrostburnS3Accessor.extend(accessor)

                        assert isinstance(accessor, fbS3.FrostburnS3Accessor), \
                            '### Re-classing {0} was unsuccessful!'

                    try:
                        # check if the resource exists
                        if accessor.exists(componentResourceIdentifier):
                            # get url to the resource
                            componentPath = accessor.getResolvedUrl(componentResourceIdentifier)

                    except ftrack.AccessorFilesystemPathError:
                        self.logger.error('### Filesystem path could not be determined from '
                                     'resourceIdentifier for component: {0}!!!\n{1}'
                                     .format(component, traceback.print_exc()))
                        return

                    except ftrack.AccessorUnsupportedOperationError:
                        self.logger.error('### Retrieving filesystem paths is not supported by '
                                     'this accessor for component: {0}!!!\n{1}'
                                     .format(component, traceback.print_exc()))
                        return

                # store data in dict
                componentData[componentName] = componentPath

            self.logger.info('Data on AssetVersion: {0}'.format(componentData))
            return componentData


def async(fn):
    """
    This decorator method allows for running *fn* asynchronously.
    :param fn:
    :return:
    """

    def wrapper(*args, **kwargs):
        thread = threading.Thread(target=fn, args=args, kwargs=kwargs)
        thread.start()
    return wrapper


def downloadFile(url, saveDirectory, fileName=None):
    """
    This method downloads a single file given a url.
    """

    logger = logging.getLogger(__name__)

    # Download the file from the URL and get fileHandle to object
    try: remoteFile = urllib2.urlopen(url)
    except urllib2.HTTPError:
        logger.error('### URL: {0} could not be downloaded!\n{1}'
                          .format(url, traceback.print_exc()))
        return None

    # if no filename is provided, generate one from the URL directly
    if not fileName:

        parsedUrl = urlparse.urlparse(remoteFile.url)
        fileName = parsedUrl[2].split('/')[-1]

    savePath = os.path.join(saveDirectory, fileName)

    logger.info('Saving file to: {0}'.format(savePath))

    # Write file to disk
    try:
        with open(savePath, 'w') as localFile:
            localFile.write(remoteFile.read())

    except IOError:
        logger.error('### Could not write to" {0}!\n{1}'.format(savePath, traceback.print_exc()))
        raise IOError

    return localFile.name

def initializeDownloads(urls, saveDirectory, fileNames=None):

    # Make list for appending all threads to
    threads = []

    # Iterate over each URL and start a thread to handle each download
    for idx, url in enumerate(urls):

        # Check if list of specific fileNames are given as well
        if fileNames:
            fileName = fileNames[idx]

        else:
            fileName = None

        # Construct thread and pass arguments
        thread = threading.Thread(
            target=downloadFile,
            kwargs={
              'url': url,
              'saveDirectory': saveDirectory,
              'fileName': fileName
        })

        threads.append(thread)
        thread.start()


def register(registry, **kwargs):
    """
    Register action. Called when used as an event plugin.
    :param registry:
    :param kwargs:
    :return:
    """

    logger = logging.getLogger(__name__)

    logger.debug('Registering action: {0}'.format(
        os.path.abspath(os.path.dirname(__file__))))

    action = DownloadAllComponentsFromEntityAction()
    action.register()


def main(arguments=None):
    """
    This entry point sets up logging and registers action when this is run
    as a standalone action.
    :param arguments:
    :return:
    """
    if arguments is None:
        arguments = []

    parser = argparse.ArgumentParser()
    # Allow setting of logging level from arguments.
    loggingLevels = {}
    for level in (
        logging.NOTSET, logging.DEBUG, logging.INFO, logging.WARNING,
        logging.ERROR, logging.CRITICAL
    ):
        loggingLevels[logging.getLevelName(level).lower()] = level

    parser.add_argument(
        '-v', '--verbosity',
        help='Set the logging output verbosity.',
        choices=loggingLevels.keys(),
        default='info'
    )
    namespace = parser.parse_args(arguments)

    # Set up basic logging
    logging.basicConfig(level=loggingLevels[namespace.verbosity])

    # Subscribe to action.
    ftrack.setup()
    action = DownloadAllComponentsFromEntityAction()
    action.register()

    # Wait for events
    ftrack.EVENT_HUB.wait()


if __name__ == '__main__':
    raise SystemExit(main(sys.argv[1:]))