load_shaders_from_ass.py

from arnold import *
import GafferArnold
import ctypes
import imath
import Gaffer

def convert_parameters( gaffer_shader, node ):
	
	node_name = AiNodeGetName( node )
	entry = AiNodeGetNodeEntry( node )

	it  = AiNodeEntryGetParamIterator( entry )
		
	while not AiParamIteratorFinished(it):

		paramEntry = AiParamIteratorGetNext(it)
		param_name  = AiParamGetName( paramEntry )
		param_type = AiParamGetType ( paramEntry )
		is_linked =  AiNodeIsLinked(node, param_name)


		if is_linked:
			component = ctypes.c_int()
	
			source_linked_node  = AiNodeGetLink( node, param_name, ctypes.byref(component) )
		
			if source_linked_node:

				linked_child_shader = create_shader_node_recursive( source_linked_node )
				if linked_child_shader:

					if component.value == -1:
						try:
							gaffer_shader["parameters"][param_name].setInput( linked_child_shader["out"] )
						except Exception,e:
							try:
								gaffer_shader["parameters"][param_name]["r"].setInput( linked_child_shader["out"]["r"] )
								gaffer_shader["parameters"][param_name]["g"].setInput( linked_child_shader["out"]["g"] )
								gaffer_shader["parameters"][param_name]["b"].setInput( linked_child_shader["out"]["b"] )
								Gaffer.Metadata.registerValue( linked_child_shader["out"], 'compoundNumericNodule:childrenVisible', True )
								Gaffer.Metadata.registerValue( gaffer_shader["parameters"][param_name], 'compoundNumericNodule:childrenVisible', True )
							except Exception,e :
								pass

					if component.value == 0: 
						Gaffer.Metadata.registerValue( linked_child_shader["out"], 'compoundNumericNodule:childrenVisible', True )
						Gaffer.Metadata.registerValue( gaffer_shader["parameters"][param_name], 'compoundNumericNodule:childrenVisible', True )
						gaffer_shader["parameters"][param_name].setInput( linked_child_shader["out"]["r"] )
					if component.value == 1:
						Gaffer.Metadata.registerValue( linked_child_shader["out"], 'compoundNumericNodule:childrenVisible', True )
						Gaffer.Metadata.registerValue( gaffer_shader["parameters"][param_name], 'compoundNumericNodule:childrenVisible', True )
						gaffer_shader["parameters"][param_name].setInput( linked_child_shader["out"]["g"] )
					if component.value == 2: 
						Gaffer.Metadata.registerValue( linked_child_shader["out"]	, 'compoundNumericNodule:childrenVisible', True )
						Gaffer.Metadata.registerValue( gaffer_shader["parameters"][param_name], 'compoundNumericNodule:childrenVisible', True )
						gaffer_shader["parameters"][param_name].setInput( linked_child_shader["out"]["b"] )

		else:
			if param_type == AI_TYPE_BYTE:
				byte_value = AiNodeGetByte(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, byte_value)
				gaffer_shader["parameters"][param_name].setValue(byte_value)
			elif param_type == AI_TYPE_INT :
				int_value = AiNodeGetInt(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, int_value)
				gaffer_shader["parameters"][param_name].setValue(int_value)
			elif param_type == AI_TYPE_UINT :
				uint_value = AiNodeGetUInt(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, uint_value)
				gaffer_shader["parameters"][param_name].setValue(uint_value)
			elif param_type == AI_TYPE_FLOAT :
				float_value = AiNodeGetFlt(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, float_value)
				gaffer_shader["parameters"][param_name].setValue(float_value)
			elif param_type == AI_TYPE_STRING:
				string_value = AiNodeGetStr(node, param_name)
				if param_name != "name":
					print "setting {}.{} to {}".format(node_name, param_name, string_value)
					gaffer_shader["parameters"][param_name].setValue(string_value)
			elif param_type == AI_TYPE_BOOLEAN:		
				bool_value = AiNodeGetBool(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, bool_value)
				gaffer_shader["parameters"][param_name].setValue(bool_value)
			elif param_type == AI_TYPE_RGB:
				rgb_value = AiNodeGetRGB(node, param_name)
				imathvalue = imath.Color3f(rgb_value.r, rgb_value.g, rgb_value.b)
				print "setting {}.{} to {}".format(node_name, param_name, imathvalue)
				gaffer_shader["parameters"][param_name].setValue(imathvalue)
			elif param_type == AI_TYPE_RGBA:
				rgba_value = AiNodeGetRGBA(node, param_name)
				imathvalue = imath.Color4f(rgba_value.r, rgba_value.g, rgba_value.b, rgba_value.a)
				print "setting {}.{} to {}".format(node_name, param_name, imathvalue)
				gaffer_shader["parameters"][param_name].setValue(imathvalue)

			elif param_type == AI_TYPE_VECTOR:
				vector_value = AiNodeGetVec(node, param_name)
				imathvalue = imath.V3f(vector_value.x, vector_value.y, vector_value.z)
				print "setting {}.{} to {}".format(node_name, param_name, imathvalue)
				gaffer_shader["parameters"][param_name].setValue(imathvalue)
			elif param_type == AI_TYPE_VECTOR2:
				vector2_value = AiNodeGetVec2(node, param_name)
				imathvalue = imath.V2f(vector2_value.x, vector2_value.y)
				print "setting {}.{} to {}".format(node_name, param_name, imathvalue)
				gaffer_shader["parameters"][param_name].setValue(imathvalue)

			elif param_type == AI_TYPE_ENUM:
				enum_value = AiNodeGetStr(node, param_name)
				print "setting {}.{} to {}".format(node_name, param_name, enum_value)
				gaffer_shader["parameters"][param_name].setValue(enum_value)
			else:
				print param_type

def create_shader_node_recursive( node ):
	
	node_name = AiNodeGetName( node )

	entry = AiNodeGetNodeEntry( node )
	shader_name = AiNodeEntryGetName( entry )

	gaffer_node_name = node_name.replace('@', '_')

	gaffer_shader = GafferArnold.ArnoldShader( gaffer_node_name )
	gaffer_shader.loadShader( shader_name )
	print gaffer_shader['name'].getValue(), shader_name

	# if gaffer_shader['name'].getValue() != shader_name: 
	# 	print "failed to load shader : '{}' load result {}".format(shader_name, r)
	# 	return None

	root.addChild(gaffer_shader)
	
	convert_parameters(gaffer_shader, node)

	return gaffer_shader


def load_ass_shaders(ass_file, shader_name):

 
	AiBegin()
	AiMsgSetConsoleFlags(AI_LOG_ALL)

	#AiLoadPlugins('C:/solidangle/mtoadeploy/2015/shaders')
	AiASSLoad(ass_file, AI_NODE_ALL)

	node = AiNodeLookUpByName(shader_name)

	create_shader_node_recursive(node)


load_ass_shaders("/home/don/Downloads/lookdev_export.ass", "towerMat")