file_runner.py

__author__ = 'ben.hearn'

"""
- File Runner is a handy Qt UI class that helps with constant batch processing on files or directories with a twist!
- You can apply any external script you have written to process the files

How to use:
- Simply drag in the directories you want to iterate over into the directories window
- Type in your desired file format

- Drag in the script you want to process your files into the script window
- Type in the name of the function you want to run
- Choose whether or not you want the program to recursively process all sub directories

- Select your desired directories and your desired script
- Click Run. It's as easy as that!

Edit the class however you see fit  :)

Enjoy!
"""

from PyQt4 import QtCore, QtGui
import sys
import os
import ast
import importlib
import ntpath

class DirectoryFunctions(QtGui.QDialog):
	def __init__(self):
		QtGui.QDialog.__init__(self)

		self.appDataDir = os.path.join(os.getenv('APPDATA'), 'CUSTOM_DIR')
		if not os.path.exists(self.appDataDir):
			os.makedirs(self.appDataDir)
		self.inFocusWidget = None
		self.allowedScriptExts = ['py', 'mel']

		global dirIn
		self.dirIn = QtGui.QDialog()
		self.dirIn.resize(850, 450)
		self.dirIn.setWindowTitle("File Runner")
		self.dirIn.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
		width = self.dirIn.width()

		# Master display is the table view that shows our
		self.masterDisplay = QtGui.QTableWidget()
		self.masterDisplay.setColumnCount(2)
		self.masterDisplay.setHorizontalHeaderLabels(['Directories', 'Format'])
		# Setting the stretch to expand with the table
		self.masterDisplay.horizontalHeader().setResizeMode(0, QtGui.QHeaderView.Interactive)
		self.masterDisplay.horizontalHeader().resizeSection(0, width/2)
		self.masterDisplay.horizontalHeader().setResizeMode(1, QtGui.QHeaderView.Stretch)
		self.masterDisplay.verticalHeader().hide()
		self.masterDisplay.setObjectName('master_display')
		self.masterDisplay.itemPressed .connect(lambda: self.setFocusWidget(self.masterDisplay))

		# Script layout
		self.scriptPaths = QtGui.QTableWidget()
		self.scriptPaths.setColumnCount(3)
		self.scriptPaths.setHorizontalHeaderLabels(['Script Name', 'Function Name', 'Recursive'])
		self.scriptPaths.horizontalHeader().setResizeMode(0, QtGui.QHeaderView.Interactive)
		self.scriptPaths.horizontalHeader().resizeSection(0, width/2)
		self.scriptPaths.horizontalHeader().setResizeMode(1, QtGui.QHeaderView.Stretch)
		self.scriptPaths.horizontalHeader().setResizeMode(2, QtGui.QHeaderView.Stretch)
		self.scriptPaths.setSelectionMode(QtGui.QAbstractItemView.SingleSelection)
		self.scriptPaths.verticalHeader().hide()
		self.scriptPaths.setObjectName('script_paths')
		self.scriptPaths.itemPressed .connect(lambda: self.setFocusWidget(self.scriptPaths))

		# Create target layout
		btnRun = QtGui.QPushButton('Run On Selected')
		targetLayout = QtGui.QHBoxLayout()
		targetLayout.addWidget(btnRun)

		# Create appdata button layout
		btnCreateAppData = QtGui.QPushButton('Save Locations')
		btnRemoveDir = QtGui.QPushButton('Remove Entry')
		appLayout = QtGui.QHBoxLayout()
		appLayout.addWidget(btnCreateAppData)
		appLayout.addWidget(btnRemoveDir)

		masterLayout = QtGui.QVBoxLayout()
		masterLayout.addWidget(self.masterDisplay)
		masterLayout.addWidget(self.scriptPaths)
		masterLayout.addLayout(targetLayout)
		masterLayout.addLayout(appLayout)

		self.masterDisplay.setAcceptDrops(True)
		self.masterDisplay.dragEnterEvent = self.dragEnterEvent
		self.masterDisplay.dragMoveEvent = self.dragMoveEvent
		self.masterDisplay.dropEvent = self.dropEvent
		#self.masterDisplay.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
		#self.masterDisplay.customContextMenuRequested.connect(self.openMenu)

		self.scriptPaths.setAcceptDrops(True)
		self.scriptPaths.dragEnterEvent = self.dragEnterEvent
		self.scriptPaths.dragMoveEvent = self.dragMoveEvent
		self.scriptPaths.dropEvent = self.scriptDropEvent

		btnCreateAppData.pressed.connect(self.createAppData)
		btnRemoveDir.pressed.connect(self.removeDir)
		btnRun.pressed.connect(self.runOnSelected)

		self.dirIn.setLayout(masterLayout)
		self.dirIn.show()
		self.setupUI()

# ----------------------------------------------------------------------------------------------- #
	def setFocusWidget(self, table):
		self.inFocusWidget = table

	""" Overidden QT drag/drop events """
	def dragEnterEvent(self, event):
		event.accept()

	def dragMoveEvent(self, event):
		event.accept()

	def dropEvent(self, event):
		md = event.mimeData()
		if md.hasUrls():
			for url in md.urls():
				urlPath = str(url.path())
				if urlPath.startswith('/'):
					urlPath = urlPath[1:]
				if not os.path.isdir(urlPath):
					event.ignore()
					continue
				self.masterDisplay.insertRow(self.masterDisplay.rowCount())
				rowNum = self.masterDisplay.rowCount()-1
				item = QtGui.QTableWidgetItem(urlPath)
				formatItem = QtGui.QTableWidgetItem('')
				item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
				formatItem.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable)
				self.masterDisplay.setItem(rowNum, 0, item)
		else:
			event.ignore()
		event.accept()

# ----------------------------------------------------------------------------------------------- #

	def scriptDropEvent(self, event):
		md = event.mimeData()
		if md.hasUrls():
			for url in md.urls():
				urlPath = str(url.path())
				if urlPath.startswith('/'):
					urlPath = urlPath[1:]
				if os.path.isdir(urlPath):
					event.ignore()
					continue
				if not urlPath.split('.')[-1] in self.allowedScriptExts:
					continue
				self.scriptPaths.insertRow(self.scriptPaths.rowCount())
				rowNum = self.scriptPaths.rowCount()-1
				item = QtGui.QTableWidgetItem(urlPath)
				functionNameItem = QtGui.QTableWidgetItem('')
				item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
				functionNameItem.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable)
				self.scriptPaths.setItem(rowNum, 0, item)

				pWidget = self.setCenterCheckbox()
				self.scriptPaths.setCellWidget(rowNum, 2, pWidget)
		else:
			event.ignore()
		event.accept()

	def setCenterCheckbox(self, checked=False):
		""" Sets a centred checkbox """
		recursiveCheckBox = QtGui.QCheckBox()
		recursiveCheckBox.setChecked(checked)

		pWidget = QtGui.QWidget()
		pLayout = QtGui.QHBoxLayout()
		pLayout.addWidget(recursiveCheckBox)
		pLayout.setAlignment(QtCore.Qt.AlignCenter)
		pLayout.setContentsMargins(0,0,0,0)

		pWidget.setLayout(pLayout)
		return pWidget

# ----------------------------------------------------------------------------------------------- #

	def unpackCheckBox(self, table, row, column):
		return table.cellWidget(row, column).layout().itemAt(0).widget().isChecked()

# ----------------------------------------------------------------------------------------------- #

	def chooseDir(self):
		""" Allows the user to choose a new directory when you right click """
		dirPath = self.openWindowsBrowser()
		indices = self.masterDisplay.selectedIndexes()
		if dirPath:
			dirPath = self.fixPathing(dirPath)
			if not os.path.isdir(dirPath):
				return
			for i in indices:
				row = i.row()
				self.createQtContent(dirPath, row, 0, True)

# ----------------------------------------------------------------------------------------------- #

	def addFormat(self, formatString=''):
		""" Adds the format string to the format column """
		indices = self.masterDisplay.selectedIndexes()
		for i in indices:
			row = i.row()
			column = i.column()
			self.addFormatItem(row, column, formatString)

# ----------------------------------------------------------------------------------------------- #

	def addFormatItem(self, row, column, exportFormat):
		""" Adds a Qt item to the  """
		formatItem = QtGui.QTableWidgetItem(exportFormat)
		formatItem.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
		self.masterDisplay.setItem(row, column, formatItem)

# ----------------------------------------------------------------------------------------------- #

	def openWindowsBrowser(self):
		dirPath = QtGui.QFileDialog.getExistingDirectory(None, 'Select Directory')
		return dirPath

# ----------------------------------------------------------------------------------------------- #

	def createAppData(self):
		""" We create an appdata file in this function to store the source level directories and export location """
		appDataFile = os.path.join(self.appDataDir, 'custom_dirs.txt')
		appDataDict = {}
		scriptDataDict = {}
		tableWidgets = {self.masterDisplay:appDataDict, self.scriptPaths:scriptDataDict}
		for table in tableWidgets:
			for r in range(table.rowCount()):
				for c in range(table.columnCount()):
					if c == 0:
						source = str(table.item(r, c).text())
					elif c == 1:
						if table.item(r, c) is not None:
							format_function = str(table.item(r, c).text())
						else:
							format_function = ''
					elif c == 2:
						checkBoxResult = self.unpackCheckBox(table, r, c)
						format_function += '**' + str(checkBoxResult)

				tableWidgets[table].update({source:format_function})

		with open(appDataFile, 'w+') as appDataFile:
			appDataFile.write(str(appDataDict)+'\n')
			appDataFile.write(str(scriptDataDict))

# ----------------------------------------------------------------------------------------------- #

	def setupUI(self):
		""" Sets up the table UI on first boot """
		appDataFile = os.path.join(self.appDataDir, 'custom_dirs.txt')
		exportDirs = None
		scriptPaths = None
		if not os.path.exists(appDataFile):
			return
		else:
			with open(appDataFile) as f:
				i = 0
				for line in f:
					if i == 0:
						exportDirs = ast.literal_eval(line)
					elif i == 1:
						scriptPaths = ast.literal_eval(line)
					i += 1
			dicts = [exportDirs, scriptPaths]
			for index, d in enumerate(dicts):
				if index == 0:
					table = self.masterDisplay
				elif index == 1:
					table = self.scriptPaths
				if len(d) > 0:
					for index, source in enumerate(d):
						table.insertRow(index)
						content = d[source]

						# If our table is script_paths we need to split our value to get our true/false checkbox value
						if table.objectName() == 'script_paths':
							temp = d[source].split('**')
							content = temp[0]
							checked = temp[1]
							if checked == 'True':
								checked = True
							else:
								checked = False
							pWidget = self.setCenterCheckbox(checked)
							table.setCellWidget(index, 2, pWidget)

						self.createQtContent(source, index, 0, table)
						self.createQtContent(content, index, 1, table)

# ----------------------------------------------------------------------------------------------- #

	def createQtContent(self, content, row, column, table):
		""" Creates a QTableWidgetItem """
		item = QtGui.QTableWidgetItem(content)
		if column == 1:
			item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable)
		else:
			item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
		table.setItem(row, column, item)

# ----------------------------------------------------------------------------------------------- #

	def fixPathing(self, filePath):
		""" Replaces pathing slashes """
		return filePath.replace('\\', '/')

# ----------------------------------------------------------------------------------------------- #

	def removeDir(self):
		""" Removes the directory from the UI and saves the file
		    Important tip is to remove rows incrementally in reverse """
		try:
			indices = self.inFocusWidget.selectedIndexes()
		except AttributeError:
			QtGui.QMessageBox.information(self.masterDisplay, 'Info','Please ensure widget is active')
			return
		# Get all row index
		indexes = []
		for i in indices:
			indexes.append(i.row())
		# Reverse sort rows indexes
		indexes = sorted(indexes, reverse=True)

		# Delete rows
		for rowidx in indexes:
			print 'removing', rowidx
			self.inFocusWidget.removeRow(rowidx)

		# ----------------------------------------------------------------------------------------------- #

	def runOnSelected(self):
		""" Runs over the selected directories """
		indices = self.masterDisplay.selectedIndexes()
		scriptPathIndices   = self.scriptPaths.selectedItems()

		if len(indices) < 1 or len(scriptPathIndices)< 1:
			QtGui.QMessageBox.information(self.masterDisplay, 'Info','Please select a script and relevant directories')
			return

		row                 = scriptPathIndices[0].row()
		scriptPath          = str(self.scriptPaths.item(row, 0).text())
		scriptDir           = os.path.dirname(scriptPath)
		scriptName          = ntpath.basename(scriptPath).split('.')[0]
		functionName        = str(self.scriptPaths.item(row, 1).text())
		recursive           = self.unpackCheckBox(self.scriptPaths, row, 2)
		if not os.path.exists(scriptPath):
			QtGui.QMessageBox.information(self.masterDisplay, 'Info','Your script no longer lives in this location\nPlease remove and choose new location')
			return

		if not scriptDir in sys.path:
			sys.path.append(scriptDir)
		try:
			externalFunc = getattr(importlib.import_module(scriptName), functionName)
		except AttributeError:
			QtGui.QMessageBox.information(self.masterDisplay, 'Info','Please check script pathing and function name is correct')
			return

		print 'Running'
		for i in indices:
			row = i.row()
			sourceDir = self.masterDisplay.item(row, 0)
			formatString = self.masterDisplay.item(row, 1)

			if sourceDir is None or formatString is None or sourceDir == '' or formatString == '':
				QtGui.QMessageBox.information(self.masterDisplay, 'Info','Cannot run without a source dir or format')
				return

			sourceDir = str(sourceDir.text())
			formatString = str(formatString.text())
			for root, dirs, files in os.walk(sourceDir):
				if len(files) > 0:
					self.runExternalFunction(root, files, formatString, externalFunc)
					if recursive == False:
						break

		QtGui.QMessageBox.information(self.masterDisplay, 'Info', 'Process complete')

	def runExternalFunction(self, root, files, format, func):
		for f in files:
			if os.path.splitext(f)[-1] == format:
				myFile = os.path.join(root, f)
				func(myFile)


def run():
	app = QtGui.QApplication(sys.argv)
	ex = DirectoryFunctions()
	sys.exit(app.exec_())

run()