Quickly search and add referenced files from directory

refSearch

'''

Author: Alexander Smirnov

The script will quickly search and add referenced files from directory you will enter in the main text field
Icon buttons from the lower right corner simply the same as the native maya ones

Start will launch the 'scriptJob' that monitoring any reference scene request before opening
Stop will delete the 'scriptJob'

You can change the 'search' path without stop-start stuff if you already start the 'scriptJob'

You can change extensions in text field of the maya files and textures files that 'scriptJob' will take into consideration

In order to run the script just launch the code:

import refSearch
try:
	r.stop()
except:
	pass
r = refSearch.refSearch().UI()

'''

import maya.cmds as mc
import maya.api.OpenMaya as om2
import os
import os.path
import maya.mel as mel


class refSearch(object):

	def __init__(self):
		self.refCallBacks = []
		self.allRefFiles = {}
		self.extensions = [ 'ma', 'mb', 'jpg', 'jpeg', 'tiff', 'tga', 'png', 'psd' ]

	def UI(self):

		winName = 'ReferenceSearch'
		if mc.window( winName, ex=1 ):
			mc.deleteUI( winName )
		mainWin = mc.window( winName, s=1, t='Reference Search', mxb=0, cc=lambda *x: self.stop() )
		mainLayout = mc.formLayout( p=mainWin )

		self.pathField = mc.textField( h=25, p=mainLayout, cc=lambda *x: self.searchFiles() )
		if mc.optionVar( ex='refSearchPath' ):
			pathVar = mc.optionVar( q='refSearchPath' )
			mc.textField( self.pathField, e=1, tx=pathVar )

		self.extensionsField = mc.textField( h=25, p=mainLayout, cc=lambda *x: self.extFields() )
		if mc.optionVar( ex='refSearchExt' ):
			extVar = mc.optionVar( q='refSearchExt' )
			mc.textField( self.extensionsField, e=1, tx=extVar )
		else:
			mc.textField( self.extensionsField, e=1, tx=str(self.extensions)[1:-1].replace(' ','').replace("'",'') )
		pathButton = mc.iconTextButton( i='folder-open.png', p=mainLayout, c=self.openFolder )

		newButton = mc.iconTextButton( i='fileNew.png', p=mainLayout, c=lambda *x: self.fileCommand('new')  )
		openButton = mc.iconTextButton( i='fileOpen.png', p=mainLayout, c=lambda *x: self.fileCommand('open')  )
		saveButton = mc.iconTextButton( i='fileSave.png', p=mainLayout, c=lambda *x: self.fileCommand('save')  )

		self.startButton = mc.button( l='Start', p=mainLayout, c=lambda *x: self.start() )

		mc.formLayout( mainLayout, e=1, af=( self.pathField, "top", 10 ) )
		mc.formLayout( mainLayout, e=1, af=( self.pathField, "left", 10 ) )
		mc.formLayout( mainLayout, e=1, ac=( self.pathField, "right", 4, pathButton ) )
		mc.formLayout( mainLayout, e=1, af=( pathButton, "top", 12 ) )
		mc.formLayout( mainLayout, e=1, af=( pathButton, "right", 10 ) )

		mc.formLayout( mainLayout, e=1, ac=( self.startButton, "top", 10, self.pathField ) )
		mc.formLayout( mainLayout, e=1, af=( self.startButton, "left", 11 ) )

		mc.formLayout( mainLayout, e=1, ac=( self.extensionsField, "top", 9, self.pathField ) )
		mc.formLayout( mainLayout, e=1, ac=( self.extensionsField, "left", 10, self.startButton ) )
		mc.formLayout( mainLayout, e=1, ac=( self.extensionsField, "right", 10, newButton ) )

		mc.formLayout( mainLayout, e=1, ac=( newButton, "top", 11, pathButton ) )
		mc.formLayout( mainLayout, e=1, ac=( newButton, "right", 4, openButton ) )
		mc.formLayout( mainLayout, e=1, ac=( openButton, "top", 11, pathButton ) )
		mc.formLayout( mainLayout, e=1, ac=( openButton, "right", 4, saveButton ) )
		mc.formLayout( mainLayout, e=1, ac=( saveButton, "top", 11, pathButton ) )
		mc.formLayout( mainLayout, e=1, af=( saveButton, "right", 10 ) )

		mc.window( mainWin, e=1, w=400, h=80 )
		mc.showWindow(winName)

		self.searchFiles()

	def fileCommand(self, com):
		if com == 'new':
			mel.eval('NewScene;')
		elif com == 'open':
			mel.eval('OpenScene;')
		elif com == 'save':
			mel.eval('SaveScene;')

	def start(self):
		mc.button( self.startButton, e=1, l='Stop', c=lambda *x: self.stop() )
		refCallBack1 = om2.MSceneMessage.addCheckFileCallback(om2.MSceneMessage.kBeforeCreateReferenceCheck, self.loadRefs, clientData={})
		refCallBack2 = om2.MSceneMessage.addCallback(om2.MSceneMessage.kAfterOpen, self.loadRest )
		self.refCallBacks.append(refCallBack1)
		self.refCallBacks.append(refCallBack2)

	def searchFiles(self):
		dirPath = mc.textField( self.pathField, q=1, tx=1)
		mc.optionVar( sv=('refSearchPath', dirPath) )
		if not os.path.isdir(dirPath):
			return
		if mc.optionVar( ex='refSearchExt' ):
			extVar = mc.optionVar( q='refSearchExt' )
			self.extensions = extVar.split(',')

		for root, dirs, files in os.walk(dirPath.replace('/', '\\')):
			for name in files:
				if name.rsplit( '.', 1 )[-1] in self.extensions:
					self.allRefFiles[os.path.join(root, name)] = name

	def loadRefs(self, fileObject, clientData):
		try:
			for k, v in self.allRefFiles.iteritems():
				print k
				if fileObject.rawName() in v:
					fileObject.setRawFullName( k )
					return True
		except:
			mc.warning("Loading reference failed.")

	def loadRest(self, *args):
		types = [ ('file', 'ftn') ]
		for type, attr in types:
			nodes = mc.ls(typ=type)
			for f in nodes:
				path = mc.getAttr(f + '.' + attr).rsplit('/', 1)[-1]
				for k, v in self.allRefFiles.iteritems():
					if path in v:
						mc.setAttr(f + '.' + attr, k, type='string')

	def stop(self):
		mc.button( self.startButton, e=1, l='Start', c=lambda *x: self.start() )
		if self.refCallBacks:
			for r in self.refCallBacks:
				try:
					om2.MMessage.removeCallback(r)
				except:
					pass
			self.refCallBacks = []

	def openFolder(self):
		dirPath = mc.textField( self.pathField, q=1, tx=1)
		fileName = mc.fileDialog2(fm=2, dir=dirPath, ds=mc.optionVar( q='FileDialogStyle' ))
		if fileName:
			mc.textField( self.pathField, e=1, tx=fileName[0] )
			self.searchFiles()

	def extFields(self):
		ext = mc.textField( self.extensionsField, q=1, tx=1)
		mc.optionVar( sv=('refSearchExt', ext) )
		self.searchFiles()