Metadata Script

editMetaData.py

"""
Author: Carlo Cherisier
Date: 04.16.15
Script: MetadataUI


To run this script:
import editMetaData
editMetaData.main()
"""

from PyQt4 import QtCore
from PyQt4 import QtGui
import sys
import os
import re
import OpenImageIO
import time

## The Logging module is a way to debug your cod
import logging
logReport = logging.getLogger(__name__)
logReport.setLevel(logging.WARNING)
logging.basicConfig()

class MetadataUI(QtGui.QDialog):
    def __init__(self, *args, **kws):
        parent = kws.get('parent', None)
        super(MetadataUI, self).__init__(parent=parent)

        ## Class variables
        self.logic = Logic()
        self.folderfilepath = ''
        self.current_output_filelist = []
        self.exr_filedict = {}
        ## Main path folder qtree view
        self.maintreepath = 'X:/'

        ## Start
        self.create_window()
        self.setup_connections()
        self.populate_foldersearch_tree()

    def create_groupbox(self, name, parent, mode=None, width=None, fsize=0, returnbox=False):
        """
        Create groupBox widet
        """
        groupbox_widget = QtGui.QGroupBox(name)
        groupbox_widget.setContentsMargins(4, 19, 2, 2)
        groupbox_widget.setAlignment(4)
        box_layout = QtGui.QVBoxLayout()
        groupbox_widget.setLayout(box_layout)
        parent.addWidget(groupbox_widget)

        if width:
            groupbox_widget.setFixedWidth(width)

        if mode is not None:
            if mode == 'qlist':
                qlist = QtGui.QListWidget()

            fontsize = QtGui.QFont()
            fontsize.setPointSize(fsize)
            qlist.setFont(fontsize)
            box_layout.addWidget(qlist)

            if returnbox:
                return qlist, box_layout

            else:
                return qlist

        else:
            return box_layout, groupbox_widget

    def create_window(self):
        """
        Purpose: Contains all window elements
        """
        ## Font Variable
        fontsize = QtGui.QFont()

        ## Window Title
        self.setWindowTitle('Edit Meta Data')

        ## Main Layout
        mainLayout = QtGui.QHBoxLayout()
        self.setLayout(mainLayout)

        ## FIRST COLUMN <-------------
        layout01 = QtGui.QVBoxLayout()
        mainLayout.addLayout(layout01)

        ## Folder View Section
        box_layout = self.create_groupbox(name='Folder Search', parent=layout01, width=250, fsize=12)[0]

        ## Create QTreeView to search through folder
        self._foldersearch_tree = QtGui.QTreeView()
        box_layout.addWidget(self._foldersearch_tree)

        ## Create line to show folder path
        self._folderpath_line = QtGui.QLineEdit()

        ## Change font size for line edit wid
        fontsize = QtGui.QFont()
        fontsize.setPointSize(8)
        self._folderpath_line.setFont(fontsize)

        ## Make line edit wid uneditable
        self._folderpath_line.setReadOnly(True)
        box_layout.addWidget(self._folderpath_line)

        ## SECOND COLUMN <-------------
        ## File List Section
        self._filelist_wid = self.create_groupbox(name='File List', parent=layout01, mode='qlist', width=250, fsize=12)

        ## THIRD COLUMN <-------------
        layout03 = QtGui.QVBoxLayout()
        mainLayout.addLayout(layout03)

        ## Transfer Buttons Section
        box_layout = self.create_groupbox(name='Transfer', parent=layout03, width=80)[0]

        space = QtGui.QSpacerItem(10, 10, QtGui.QSizePolicy.Minimum, QtGui.QSizePolicy.Expanding)

        ## Add Buttons
        self._add_shots_button = QtGui.QPushButton('>')
        self._add_all_shots_button = QtGui.QPushButton('>>')

        ## Remove Buttons
        self._remove_shots_button = QtGui.QPushButton('<')
        self._remove_all_shots_button = QtGui.QPushButton('<<')

        ## Adding buttons to group box
        box_layout.addItem(space)
        box_layout.addWidget(self._add_shots_button)
        box_layout.addWidget(self._add_all_shots_button)
        box_layout.addItem(space)
        box_layout.addWidget(self._remove_shots_button)
        box_layout.addWidget(self._remove_all_shots_button)

        ## FOURTH COLUMN <-------------
        layout04 = QtGui.QVBoxLayout()
        # layout04.addStretch(1)
        mainLayout.addLayout(layout04)

        ## Output List Section
        self._outputlist_wid = self.create_groupbox(name='Output List', parent=layout04, mode='qlist', width=450, fsize=8)

         ## FIFTH COLUMN <-------------
        ## Meta Data Section
        self._metadatalist_wid, box_layout = self.create_groupbox('MetaData List', parent=mainLayout, mode='qlist', width=200, fsize=12, returnbox=True)

        ## Metada Key Value Section
        gridLayout = QtGui.QGridLayout()
        box_layout.addLayout(gridLayout)

        gridLayout.addWidget(QtGui.QLabel("Key"), 0, 0)
        gridLayout.addWidget(QtGui.QLabel("Value"), 0, 1)
        self._line_row01_col00 = QtGui.QLineEdit()
        self._line_row01_col01 = QtGui.QLineEdit()
        gridLayout.addWidget(self._line_row01_col00, 1, 0)
        gridLayout.addWidget(self._line_row01_col01, 1, 1)

        self._line_row02_col00 = QtGui.QLineEdit()
        self._line_row02_col01 = QtGui.QLineEdit()
        gridLayout.addWidget(self._line_row02_col00, 2, 0)
        gridLayout.addWidget(self._line_row02_col01, 2, 1)

        self._line_row03_col00 = QtGui.QLineEdit()
        self._line_row03_col01 = QtGui.QLineEdit()
        gridLayout.addWidget(self._line_row03_col00, 3, 0)
        gridLayout.addWidget(self._line_row03_col01, 3, 1)

        self._line_row04_col00 = QtGui.QLineEdit()
        self._line_row04_col01 = QtGui.QLineEdit()
        gridLayout.addWidget(self._line_row04_col00, 4, 0)
        gridLayout.addWidget(self._line_row04_col01, 4, 1)

        ## Add Metadata Button
        self._add_metadata_button = QtGui.QPushButton('Add Metadata To Files')
        self._add_metadata_button.setStatusTip('Add Metadata to selected files')
        box_layout.addWidget(self._add_metadata_button)

    def setup_connections(self):
        """
        Purpose: Connect all GUI Elemenents to functions
        """
        ## Extra QList Attributes
        self._filelist_wid.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
        self._outputlist_wid.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)

        ## Folder Tree
        self._foldersearch_tree.clicked.connect(self.update_foldersearch_tree)

        ## Add Buttons
        self._add_shots_button.clicked.connect(lambda: self.add_outputlist_items('selected'))
        self._add_all_shots_button.clicked.connect(lambda: self.add_outputlist_items('all'))

        ## Remove Buttons
        self._remove_shots_button.clicked.connect(lambda: self.remove_outputlist_items('selected'))
        self._remove_all_shots_button.clicked.connect(lambda: self.remove_outputlist_items('all'))

        ## Option Buttons
        self._add_metadata_button.clicked.connect(self.update_add_metadata_button)

    def populate_foldersearch_tree(self):
        '''
        Populate foldersearch QTreeView
        '''
        ## Create model system needed to view folder and files
        self.model = QtGui.QFileSystemModel()

        ## Set Root path to main path location
        self.model.setRootPath(self.maintreepath)

        ## Set filter to only view directories
        self.model.setFilter(QtCore.QDir.Dirs | QtCore.QDir.NoDotAndDotDot)

        ## Set QTreeView to the QFileSystemModel
        self._foldersearch_tree.setModel(self.model)
        ## Hide Size, Type, and Data Columns

        self._foldersearch_tree.setColumnHidden(1, True)
        self._foldersearch_tree.setColumnHidden(2, True)
        self._foldersearch_tree.setColumnHidden(3, True)

        ## Set default path
        self._foldersearch_tree.setRootIndex(self.model.index(self.maintreepath))

    def update_foldersearch_tree(self):
        '''
        Update file list with contents of selected
        item in folder search tree
        '''
        ## Clear out the file list wid
        self._filelist_wid.clear()

        if len(self._foldersearch_tree.selectedIndexes()) == 0:
            return

        ## Grab selected item in tree
        index = self._foldersearch_tree.selectedIndexes()[-1]

        ## Get index from QFileSystemModel using the index from the tree
        indexItem = self.model.index(index.row(), 0, index.parent())

        ## Get file path of select folder and convert to string
        self.folderfilepath = str(self.model.filePath(indexItem))

        ## Set folder path text in folderpath line widget
        self._folderpath_line.setText(self.folderfilepath)
        # print self.folderfilepath

        ## Use the Logic function to grab all the exr in the folder
        self.exr_filedict = self.logic.getFolderContents(self.folderfilepath)

        ## The or [] is incase exr_filedict == None this prevents it
        ## from raising an error
        for each in self.exr_filedict.values() or []:
            self._filelist_wid.addItem(each)

    def add_outputlist_items(self, mode):
        '''
        Transfer selected items from File List to Output List
        '''
        if len(self._filelist_wid.selectedItems()) == 0:
            return

        ## Grab window instance
        qApp = QtGui.QApplication.instance()

        ## Switch Normal Cursor to Loading Cursor
        qApp.setOverrideCursor(QtCore.Qt.WaitCursor)

        if mode == 'selected':
            ## Grab selected shots
            ## Convert qlist to text list
            filelist = [str(x.text()) for x in self._filelist_wid.selectedItems()]

        if mode == 'all':
            ## Grab all shots
            filelist = [str(self._filelist_wid.item(i).text()) for i in range(self._filelist_wid.count())]

        ## Folder name
        for each in filelist:
            ## Loop through dictionary to find the full file path
            ## the selected item
            for key, value in self.exr_filedict.items():
                if each == value:
                    filepath = key
                    break

            ## Avoid adding duplicates
            if filepath in str(self.current_output_filelist):
                continue

            ## Add file path to output list
            self._outputlist_wid.addItem(filepath)

            ## Store file in list for later use
            self.current_output_filelist.append(filepath)

        ## Return Cursor to Normal Mode
        qApp.restoreOverrideCursor()

    def remove_outputlist_items(self, mode):
        """
        Remove shot files from Update List Widget
        """
        if mode == 'selected':
            ## Grab selected items and convet them to str
            remove_shotlist = [str(x.text()) for x in self._outputlist_wid.selectedItems()]

            ## Remove all items from qlist widget
            self._outputlist_wid.clear()

            for each in self.current_updatelist:
                if each not in remove_shotlist:
                    self._outputlist_wid.addItem(each)

                else:
                    ## Remove item from list
                    self.current_output_filelist.remove(each)

        if mode == 'all':
            ## Remove all items from qlist widget
            self._outputlist_wid.clear()

            ## Remove all items from list
            self.current_output_filelist = []

    def update_add_metadata_button(self):
        '''
        Grab key and values and add them to
        files in output file list
        '''
        ## I output file list is empty to add metadata
        if not self.current_output_filelist:
            return

        ## Create dictionary to store key and values
        ## from ui
        data_dict = {}
        for i in range(1, 5):
            cmd = 'str(self._line_row0{0}_col00.text())'.format(i)
            key = eval(cmd)

            cmd = 'str(self._line_row0{0}_col01.text())'.format(i)
            value = eval(cmd)

            if key:
                data_dict[key] = value

        ## Edit Metadata
        self.logic.add_metadata(self.current_output_filelist, data_dict)

#####################################################
class Logic(object):
    '''
    This class will contain all of the function that
    do not deal directly with the GUI
    You should always keep them separate so that it is easier
    to debug your script
    '''
    def getFolderContents(self, filepath, searchfor='exr'):
        '''
        Grab contents of folder
        searchfor works as a filter
        '''
        ## Create dictionary to hold exr file path
        exr_filedict = {}

        for root, dirs, files in os.walk(filepath):
            ## This is called list comprehension
            ## it's a simpler of creating a list
            cmd = '.*\.{0}$'.format(searchfor)
            exr_Files = [os.path.join(root, f) for f in files if re.match(cmd, f)]

            for each in exr_Files:
                ## Split off file name from path and store in dictionary
                each = each.replace( '\\', '/')
                exr_filedict[each] = os.path.split(each)[-1]

        # exr_filelist = [x for x in os.listdir(filepath) if searchfor in x]
        logReport.debug('\n'.join(exr_filedict.keys()))

        return exr_filedict

    def add_metadata(self, filepath_list=[], data_dict={}, filetype = 'exr'):
        '''
        Add medata for each file in file list
        '''
        ## Loops through file list
        for each in filepath_list:
            ## Loops through dictionary
            for key, value in data_dict.items():
                ## Renames Channels
                newimg = OpenImageIO.ImageBuf(each)

                ## Adds metadata
                newimg.specmod().attribute(key, value)

                ## Create new file path
                tmpPath = re.sub( r'(?=.*)(?=/\w*\.{0})'.format(filetype),'_tmp', path)
                
                ## Writes image
                newimg.write(tmpPath)

                ## Pause
                time.sleep(.01)

                ## Overwrites image 
                img = OpenImageIO.ImageBuf(tmpPath)
                img.write(each)

                ## Pause
                time.sleep(.01)

                ## Delete tmp path
                os.remove(tmpPath)

                logReport.info('Adding this key:{0} and value:{1} to file:{2}'.format(key, value, each))

        logReport.info('FINISHED')



def show_dialog():
    """
    Shows the dialog as a singleton
    """
    ## Make variable global to be able to check for it later
    global g_AssetUpdateUI_instance

    try:
        ## Check if UI instance is present if so close it
        ## So that there ain't two UI in the same session
        g_AssetUpdateUI_instance.closeAndDeleteLater()
    except:
        pass

    ## Store UI in global variable
    g_AssetUpdateUI_instance = MetadataUI()
    g_AssetUpdateUI_instance.show()

    ## Return variable so that it is live in the session
    return g_AssetUpdateUI_instance


def main():
    app = QtGui.QApplication(sys.argv)
    show_dialog()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()