scripted/slopeShader.py

#-
# ==========================================================================
# Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
#+

# import maya
# maya.cmds.loadPlugin("slopeShader.py")
#
# Create a sphere and assign a shader such as lambert to it.
# Select the color channel and link it to Utilities/Sp SlopeShader
# Render the scene

# imports
import maya.OpenMaya as OpenMaya
import maya.OpenMayaUI as OpenMayaUI
import maya.OpenMayaMPx as OpenMayaMPx
import math, sys

# consts
kSlopeShaderBehaviourName = "slopeShaderBehaviour"
kSlopeNodeName = "spSlopeShader"
kSlopeNodeClassify = "utility/color"
kSlopeNodeId = OpenMaya.MTypeId(0x87001)

# Node definition
class slopeShader(OpenMayaMPx.MPxNode):
    # class variables
    aAngle = OpenMaya.MObject()
    aColor1 = OpenMaya.MObject()
    aColor2 = OpenMaya.MObject()
    aTriangleNormalCamera = OpenMaya.MObject()
    aOutColor = OpenMaya.MObject()
    aMatrixEyeToWorld = OpenMaya.MObject()
    aDirtyShaderAttr = OpenMaya.MObject()

    def __init__(self):
        OpenMayaMPx.MPxNode.__init__(self)


    def compute(self, plug, dataBlock):
        if plug == slopeShader.aOutColor or plug.parent() == slopeShader.aOutColor:

            resultColor = OpenMaya.MFloatVector(0.0,0.0,0.0)

            try:
                dataHandle = dataBlock.inputValue( slopeShader.aColor1 )
            except:
                sys.stderr.write( "Failed to get inputValue aColor1" )
                raise
            walkable = dataHandle.asFloatVector()

            try:
                dataHandle = dataBlock.inputValue( slopeShader.aColor2 )
            except:
                sys.stderr.write( "Failed to get inputValue aColor2" )
                raise
            nonWalkable = dataHandle.asFloatVector()

            try:
                dataHandle = dataBlock.inputValue( slopeShader.aTriangleNormalCamera )
            except:
                sys.stderr.write( "Failed to get inputValue aTriangleNormalCamera" )
                raise
            surfaceNormal = dataHandle.asFloatVector()

            try:
                dataHandle = dataBlock.inputValue( slopeShader.aMatrixEyeToWorld )
            except:
                sys.stderr.write( "Failed to get inputValue aMatrixEyeToWorld" )
                raise
            viewMatrix = dataHandle.asFloatMatrix()

            # Get the angle
            try:
                dataHandle = dataBlock.inputValue( slopeShader.aAngle )
            except:
                sys.stderr.write( "Failed to get inputValue aAngle" )
                raise
            angle = dataHandle.asFloat()

            yVector = OpenMaya.MFloatVector(0, 1, 0)

            # Normalize the view vector
            #
            surfaceNormal.normalize()
            WSVector = surfaceNormal * viewMatrix

            # find dot product
            #
            scalarNormal = WSVector * yVector

            # take the absolute value
            #
            if scalarNormal < 0.0:
                scalarNormal *= -1.0

            radianAngle = math.radians(angle)
            cosOfAngle = math.cos(radianAngle)
            if cosOfAngle < scalarNormal:
                resultColor = walkable
            else:
                resultColor = nonWalkable

            # set output color attribute
            #
            try:
                outColorHandle = dataBlock.outputValue( slopeShader.aOutColor )
            except:
                sys.stderr.write( "Failed to get outputValue aOutColor" )
                raise

            outColorHandle.setMFloatVector(resultColor)
            outColorHandle.setClean()

        else:
            return OpenMaya.kUnknownParameter


    def postConstructor(self):
        print "In slopeShader.postConstructor"
        # OpenMayaMPx.MPxNode.setMPSafe(self, 1)


class slopeShaderBehavior(OpenMayaMPx.MPxDragAndDropBehavior):
    def __init__(self):
        OpenMayaMPx.MPxDragAndDropBehavior.__init__(self)


    def shouldBeUsedFor(self, sourceNode, destinationNode, sourcePlug, destinationPlug):
        """
        Overloaded function from MPxDragAndDropBehavior
        this method will return True if it is going to handle the connection
        between the two nodes given.
        """
        result = False

        if sourceNode.hasFn(OpenMaya.MFn.kLambert):
            #if the source node was a lambert
            #than we will check the downstream connections to see 
            #if a slope shader is assigned to it.
            #
            shaderNode = None
            src = OpenMaya.MFnDependencyNode(sourceNode)
            connections = OpenMaya.MPlugArray()
            src.getConnections(connections)

            for i in range(connections.length()):
                #check the incoming connections to this plug
                #
                connectedPlugs = OpenMaya.MPlugArray()
                connections[i].connectedTo(connectedPlugs, True, False)
                for j in range(connectedPlugs.length()):
                    #if the incoming node is a slope shader than 
                    #set shaderNode equal to it and break out of the inner 
                    #loop
                    #
                    testNode = OpenMaya.MFnDependencyNode(connectedPlugs[j].node())
                    if testNode.typeName() == kSlopeNodeName:
                        shaderNode = connectedPlugs[j].node()
                        break

                #if the shaderNode is not None
                #than we have found a slopeShader
                #
                if shaderNode is not None:
                    #if the destination node is a mesh than we will take
                    #care of this connection so set the result to True
                    #and break out of the outer loop
                    #
                    if destinationNode.hasFn(OpenMaya.MFn.kMesh):
                        result = True
                    break

        node = OpenMaya.MFnDependencyNode(sourceNode)
        if node.typeName() == kSlopeNodeName:
            #if the sourceNode is a slope shader than check what we
            #are dropping on to
            #
            if destinationNode.hasFn(OpenMaya.MFn.kLambert):
                result = True
            elif destinationNode.hasFn(OpenMaya.MFn.kMesh):
                result = True

        return result


    def connectNodeToNode(self, sourceNode, destinationNode, force):
        """
        Overloaded function from MPxDragAndDropBehavior
        this method will handle the connection between the slopeShader and the shader it is
        assigned to as well as any meshes that it is assigned to.
        """
        src = OpenMaya.MFnDependencyNode(sourceNode)

        #if we are dragging from a lambert
        #we want to check what we are dragging
        #onto.
        if sourceNode.hasFn(OpenMaya.MFn.kLambert):
            shaderNode = OpenMaya.MObject()
            connections = OpenMaya.MPlugArray()
            shaderNodes = OpenMaya.MObjectArray()
            shaderNodes.clear()

            #if the source node was a lambert
            #than we will check the downstream connections to see 
            #if a slope shader is assigned to it.
            #
            src.getConnections(connections)
            for i in range(connections.length()):
                #check the incoming connections to this plug
                #
                connectedPlugs = OpenMaya.MPlugArray()
                connections[i].connectedTo(connectedPlugs, True, False)
                for j in range(connectedPlugs.length()):
                    #if the incoming node is a slope shader than 
                    #append the node to the shaderNodes array
                    #
                    currentnode = connectedPlugs[j].node()
                    if OpenMaya.MFnDependencyNode(currentnode).typeName() == kSlopeNodeName:
                        shaderNodes.append(currentnode)

            #if we found a shading node
            #than check the destination node 
            #type to see if it is a mesh
            #
            if shaderNodes.length() > 0:
                dest = OpenMaya.MFnDependencyNode(destinationNode)
                if destinationNode.hasFn(OpenMaya.MFn.kMesh):
                    #if the node is a mesh than for each slopeShader
                    #connect the worldMesh attribute to the dirtyShaderPlug
                    #attribute to force an evaluation of the node when the mesh
                    #changes
                    #
                    for i in range(shaderNodes.length()):
                        srcPlug = dest.findPlug("worldMesh")
                        destPlug = OpenMaya.MFnDependencyNode(shaderNodes[i]).findPlug("dirtyShaderPlug")

                        if (not srcPlug.isNull()) and (not destPlug.isNull()):
                            cmd = "connectAttr -na %s %s" % (srcPlug.name(), destPlug.name())
                            OpenMaya.MGlobal.executeCommand(cmd)

                    #get the shading engine so we can assign the shader
                    #to the mesh after doing the connection
                    #
                    shadingEngine = self._findShadingEngine(sourceNode)

                    #if there is a valid shading engine than make
                    #the connection
                    #
                    if not shadingEngine.isNull():
                        cmd = "sets -edit -forceElement %s %s" % (
                                        OpenMaya.MFnDependencyNode(shadingEngine).name(),
                                        OpenMaya.MFnDagNode(destinationNode).partialPathName()
                                                                )
                        OpenMaya.MGlobal.executeCommand(cmd)

        elif src.typeName() == kSlopeNodeName:
            #if we are dragging from a slope shader
            #than we want to see what we are dragging onto
            #
            if destinationNode.hasFn(OpenMaya.MFn.kMesh):
                #if the user is dragging onto a mesh
                #than make the connection from the worldMesh
                #to the dirtyShader plug on the slopeShader
                #
                dest = OpenMaya.MFnDependencyNode(destinationNode)
                srcPlug = dest.findPlug("worldMesh")
                destPlug = src.findPlug("dirtyShaderPlug")
                if (not srcPlug.isNull()) and (not destPlug.isNull()):
                    cmd = "connectAttr -na %s %s" % (srcPlug.name(), destPlug.name())
                    OpenMaya.MGlobal.executeCommand(cmd)

            elif destinationNode.hasFn(OpenMaya.MFn.kLambert):
                dest = OpenMaya.MFnDependencyNode(destinationNode)
                srcPlug = src.findPlug("outColor")
                destPlug = dest.findPlug("color")
                if (not srcPlug.isNull()) and (not destPlug.isNull()):
                    if force:
                        cmd = "connectAttr -f %s %s" % (srcPlug.name(), destPlug.name())
                    else:
                        cmd = "connectAttr %s %s" % (srcPlug.name(), destPlug.name())
                    OpenMaya.MGlobal.executeCommand(cmd)


    def connectNodeToAttr(self, sourceNode, destinationPlug, force):
        """
        Overloaded function from MPxDragAndDropBehavior
        this method will assign the correct output from the slope shader 
        onto the given attribute.
        """
        src = OpenMaya.MFnDependencyNode(sourceNode)

        #if we are dragging from a slopeShader
        #to a shader than connect the outColor
        #plug to the plug being passed in
        #
        if destinationPlug.node().hasFn(OpenMaya.MFn.kLambert):
            if src.typeName() == kSlopeNodeName:
                srcPlug = src.findPlug("outColor")
                if (not srcPlug.isNull()) and (not destinationPlug.isNull()):
                    cmd = "connectAttr %s %s" % (srcPlug.name(), destinationPlug.name())
                    OpenMaya.MGlobal.executeCommand(cmd)
        else:
            #in all of the other cases we do not need the plug just the node
            #that it is on
            #
            destinationNode = destinationPlug.node()
            self.connectNodeToNode(sourceNode, destinationNode, force)


    def _findShadingEngine(self, node):
        """
        Given the material MObject this method will
        return the shading group that it is assigned to.
        if there is no shading group associated with
        the material than a null MObject is passed back.
        """
        nodeFn = OpenMaya.MFnDependencyNode (node)
        srcPlug = nodeFn.findPlug("outColor")
        nodeConnections = OpenMaya.MPlugArray()
        srcPlug.connectedTo(nodeConnections, False, True)
        #loop through the connections
        #and find the shading engine node that
        #it is connected to
        #
        for i in range(nodeConnections.length()):
            theNode = nodeConnections[i].node()
            if theNode.hasFn(OpenMaya.MFn.kShadingEngine):
                return theNode

        #no shading engine associated so return a
        #null MObject
        #
        return OpenMaya.MObject()


##################################################################


def nodeCreator():
    return OpenMayaMPx.asMPxPtr(slopeShader())


def behaviourCreator():
    return OpenMayaMPx.asMPxPtr(slopeShaderBehavior())


def nodeInitializer():
    nAttr = OpenMaya.MFnNumericAttribute()
    nMAttr = OpenMaya.MFnMatrixAttribute()
    nTAttr = OpenMaya.MFnTypedAttribute()
    nGAttr = OpenMaya.MFnGenericAttribute()
    # input
    slopeShader.aAngle = nAttr.create( "angle", "ang", OpenMaya.MFnNumericData.kFloat )
    nAttr.setDefault(30.0)
    nAttr.setMin(0.0)
    nAttr.setMax(100.0)
    nAttr.setKeyable(1)
    nAttr.setStorable(1)
    nAttr.setReadable(1)
    nAttr.setWritable(1)

    slopeShader.aColor1 = nAttr.createColor("walkableColor", "w")
    nAttr.setDefault(0.0, 1.0, 0.0)
    nAttr.setKeyable(1)
    nAttr.setStorable(1)
    nAttr.setUsedAsColor(1)
    nAttr.setReadable(1)
    nAttr.setWritable(1)

    slopeShader.aColor2 = nAttr.createColor( "nonWalkableColor", "nw" )
    nAttr.setDefault(1.0, 0.0, 0.0)
    nAttr.setKeyable(1)
    nAttr.setStorable(1)
    nAttr.setUsedAsColor(1)
    nAttr.setReadable(1)
    nAttr.setWritable(1)

    # Surface Normal supplied by the render sampler
    #
    slopeShader.aTriangleNormalCamera = nAttr.createPoint( "triangleNormalCamera", "n" )
    nAttr.setStorable(0)
    nAttr.setHidden(1)
    nAttr.setReadable(1)
    nAttr.setWritable(1)

    # View matrix from the camera into world space
    #
    slopeShader.aMatrixEyeToWorld = nMAttr.create( "matrixEyeToWorld", "mew", OpenMaya.MFnMatrixAttribute.kFloat )
    nAttr.setHidden(1)
    nMAttr.setWritable(1)

    # Output Attributes
    #
    slopeShader.aOutColor  = nAttr.createColor( "outColor", "oc" )
    nAttr.setStorable(0)
    nAttr.setHidden(0)
    nAttr.setReadable(1)
    nAttr.setWritable(0)

    meshTypeId = OpenMaya.MTypeId(OpenMaya.MFnData.kMesh)

    # dummy plug for forcing evaluation
    #
    slopeShader.aDirtyShaderAttr = nGAttr.create( "dirtyShaderPlug", "dsp")
    nGAttr.setArray(1)
    nGAttr.setHidden(0)
    nGAttr.setUsesArrayDataBuilder(1)
    nGAttr.setReadable(0)
    nGAttr.setStorable(1)
    nGAttr.setIndexMatters(0)
    # nGAttr.addAccept(meshTypeId)

    # Add attribues
    #
    slopeShader.addAttribute(slopeShader.aAngle)
    slopeShader.addAttribute(slopeShader.aColor1)
    slopeShader.addAttribute(slopeShader.aColor2)
    slopeShader.addAttribute(slopeShader.aTriangleNormalCamera)
    slopeShader.addAttribute(slopeShader.aOutColor)
    slopeShader.addAttribute(slopeShader.aMatrixEyeToWorld)
    slopeShader.addAttribute(slopeShader.aDirtyShaderAttr)

    slopeShader.attributeAffects (slopeShader.aAngle,  slopeShader.aOutColor)
    slopeShader.attributeAffects (slopeShader.aColor1, slopeShader.aOutColor)
    slopeShader.attributeAffects (slopeShader.aColor2, slopeShader.aOutColor)
    slopeShader.attributeAffects (slopeShader.aTriangleNormalCamera,  slopeShader.aOutColor)
    slopeShader.attributeAffects (slopeShader.aDirtyShaderAttr, slopeShader.aOutColor)


# initialize the script plug-in
def initializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject)

    # register node
    try:
        mplugin.registerNode(kSlopeNodeName, kSlopeNodeId, nodeCreator, nodeInitializer, OpenMayaMPx.MPxNode.kDependNode, kSlopeNodeClassify)
    except:
        sys.stderr.write("Failed to register node: %s" % kSlopeNodeName )
        raise

    # register behaviors
    try:
        mplugin.registerDragAndDropBehavior(kSlopeShaderBehaviourName, behaviourCreator)
    except:
        sys.stderr.write("Failed to register behaviour: %s" % kSlopeShaderBehaviourName)

    postCmd = "if( `window -exists createRenderNodeWindow` )  {refreshCreateRenderNodeWindow(\"%s\");}\n" % kSlopeNodeClassify
    OpenMaya.MGlobal.executeCommand(postCmd)


# uninitialize the script plug-in
def uninitializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject, "Autodesk", "1.0", "Any")
    try:
        mplugin.deregisterNode(kSlopeNodeId)
    except:
        sys.stderr.write( "Failed to unregister node: %s" % kSlopeNodeName )
        raise

    try:
        plugin.deregisterDragAndDropBehavior(kSlopeShaderBehaviourName)
    except:
        sys.stderr.write("Failed to deregister behaviour: %s" % kSlopeShaderBehaviourName)

    postCmd = "if( `window -exists createRenderNodeWindow` )  {refreshCreateRenderNodeWindow(\"%s\");}\n" % kSlopeNodeClassify
    OpenMaya.MGlobal.executeCommand(postCmd)