scripted/instanceShape.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.
#
# ==========================================================================
#+

###############################################################################
##
## instanceShape.py
##
## Description:
##    Registers a new shape that acts like an instancer.  The new shape
##    type is called "instanceShape".
##
##    The shape will instance N copies of a shape connected via a message
##    attribute on the node.  The sample will distribute these N copies
##    in the XZ plane.
##
##    There are no output attributes for this shape.
##    The following input attributes define the type of shape to draw.
##
##       radius        : circle radius for instance object. 
##       instanceShape : a connection to the shape to instance 
##       count         : number of instances to make. 
##
##    Additionally the instancing feature demonstrated in this code
##    only works for custom shapes. Non-custom shapes will not work.
##
################################################################################

# Usage:
# import maya
# maya.cmds.loadPlugin("instanceShape.py")
# maya.cmds.loadPlugin("basicShape.py")

# basicShape = maya.cmds.createNode("spBasicShape")
# instanceShape = maya.cmds.createNode("spInstanceShape")
# maya.cmds.connectAttr( basicShape + ".message", instanceShape + ".instanceShape" )
# 

import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaRender as OpenMayaRender
import maya.OpenMayaUI as OpenMayaUI

import math
import sys

kPluginNodeTypeName = "spInstanceShape"
spInstanceShapeNodeId = OpenMaya.MTypeId(0x00080029)

glRenderer = OpenMayaRender.MHardwareRenderer.theRenderer()
glFT = glRenderer.glFunctionTable()

kLeadColor              = 18 # green
kActiveColor            = 15 # white
kActiveAffectedColor    = 8  # purple
kDormantColor           = 4  # blue
kHiliteColor            = 17 # pale blue

kDefaultRadius = 1.0
kDefaultCount  = 10


#####################################################################
##
## Geometry class
##
class instanceGeom:
    def __init__(self): 
        self.radius = kDefaultRadius
        self.count = kDefaultCount
        self.instanceShape = None
        self.drawQueueList = [] 
        
#####################################################################
##
## Shape class - defines the non-UI part of a shape node
##
class instanceShape(OpenMayaMPx.MPxSurfaceShape):
    # class variables
    aRadius = OpenMaya.MObject()
    aCount = OpenMaya.MObject()
    aInstanceShape = OpenMaya.MObject()
    
    def __init__(self):
        OpenMayaMPx.MPxSurfaceShape.__init__(self)
        
        # geometry
        self.__myGeometry = instanceGeom()

    # override
    def postConstructor(self):
        """
         When instances of this node are created internally, the
         MObject associated with the instance is not created until
         after the constructor of this class is called. This means
         that no member functions of MPxSurfaceShape can be called in
         the constructor.  The postConstructor solves this
         problem. Maya will call this function after the internal
         object has been created.  As a general rule do all of your
         initialization in the postConstructor.
        """
        self.setRenderable(True)

    # override
    def getInternalValue(self, plug, datahandle):
        """
         Handle internal attributes.
         In order to impose limits on our attribute values we
         mark them internal and use the values in fGeometry instead.
        """
        if (plug == instanceShape.aRadius):
            datahandle.setDouble(self.__myGeometry.radius)
        elif (plug == instanceShape.aCount):
            datahandle.setInt(self.__myGeometry.count)
        else:
            return OpenMayaMPx.MPxSurfaceShape.getInternalValue(self, plug, datahandle)

        return True


    # override
    def setInternalValue(self, plug, datahandle):
        """
         Handle internal attributes.
         In order to impose limits on our attribute values we
         mark them internal and use the values in fGeometry instead.
        """

        # the minimum radius is 0
        #
        if (plug == instanceShape.aRadius):
            radius = datahandle.asDouble()
            
            if (radius < 0):
                radius = 0

            self.__myGeometry.radius = radius

        elif (plug == instanceShape.aCount):
            count = datahandle.asInt()
            if (count < 0):
                count = 0 
            self.__myGeometry.count = count
        else:
            return OpenMayaMPx.MPxSurfaceShape.setInternalValue(self, plug, datahandle)
        return True


    # override
    def isBounded(self):
        return True


    # override
    def boundingBox(self):
        """
         Returns the bounding box for the shape.
         In this case just use the radius and height attributes
         to determine the bounding box.
        """
        result = OpenMaya.MBoundingBox()

        geom = self.geometry()
        
        # Include the instance shape bounding box
        if geom.instanceShape: 
            fnDag = OpenMaya.MFnDagNode( geom.instanceShape ) 
            result = fnDag.boundingBox() 

        r = geom.radius
        instanceBbox = OpenMaya.MBoundingBox( result ) 
        for c in range( geom.count ): 
            percent = float(c)/float(geom.count)
            rad = 2*math.pi * percent
            p = (r*math.cos(rad), r*math.sin(rad),0.0)
            newbbox = OpenMaya.MBoundingBox( instanceBbox )
            trans = OpenMaya.MTransformationMatrix( )   
            vec = OpenMaya.MVector( p[0], p[1], p[2] ) 
            trans.setTranslation( vec, OpenMaya.MSpace.kTransform ) 
            mmatrix = trans.asMatrix(); 
            newbbox.transformUsing( mmatrix ) 
            result.expand( newbbox ) 
        
        return result


    def geometry(self):
        """
         This function gets the values of all the attributes and
         assigns them to the fGeometry. Calling MPlug::getValue
         will ensure that the values are up-to-date.
        """
        # return self.__myGeometry

        this_object = self.thisMObject()

        plug = OpenMaya.MPlug(this_object, instanceShape.aRadius)
        self.__myGeometry.radius = plug.asDouble()
        plug = OpenMaya.MPlug(this_object, instanceShape.aCount)
        self.__myGeometry.count = plug.asInt()

        plug = OpenMaya.MPlug(this_object, instanceShape.aInstanceShape)
        plugArray = OpenMaya.MPlugArray()
        plug.connectedTo( plugArray, True, False )
        if ( plugArray.length() > 0 ): 
            node = plugArray[0].node()
            dagNode = OpenMaya.MFnDagNode(node)
            path = OpenMaya.MDagPath()
            dagNode.getPath(path)
            self.__myGeometry.instanceShape = path 

        return self.__myGeometry

#####################################################################
##
## UI class - defines the UI part of a shape node
##
class instanceShapeUI(OpenMayaMPx.MPxSurfaceShapeUI):
    # private enums
    def __init__(self):
        OpenMayaMPx.MPxSurfaceShapeUI.__init__(self)

    # override
    def getDrawRequests(self, info, objectAndActiveOnly, queue):
        """
         The draw data is used to pass geometry through the 
         draw queue. The data should hold all the information
         needed to draw the shape.
        """
        self.geometry = None

        # Custom instancer objects can instance other custom surface 
        # shapes.  This is done by first getting the draw request 
        # data for the instancer shape 
        #
        data = OpenMayaUI.MDrawData()
        request = info.getPrototype(self)
        shapeNode = self.surfaceShape()
        path = info.multiPath()
        view = info.view() 
        # We stored the instance object via a connection on the surface
        # shape.  Retrieve that value.
        #
        geom = shapeNode.geometry()
        shadedMode = False
        mainMaterial = None 
        if request.displayStyle() == OpenMayaUI.M3dView.kGouraudShaded:
            shadedMode = True
            
        if geom.instanceShape:
            # Find the MPxSurfaceShape for the object that we are instancing
            #
            shapeUI = OpenMayaMPx.MPxSurfaceShapeUI.surfaceShapeUI(geom.instanceShape)
            if shapeUI:     
                mainMaterial = shapeUI.material( geom.instanceShape )
                mainMaterial.evaluateMaterial( view, geom.instanceShape ) 
                r = geom.radius 
                for a in range(geom.count):
                    myQueue = OpenMayaUI.MDrawRequestQueue()
                    percent = float(a)/float(geom.count)
                    rad = 2*math.pi * percent
                    position = (r*math.cos(rad), r*math.sin(rad),0.0)
                    # Construct a reference to MDrawInfo and modify it 
                    # to point to the instance shape.  If we do not do 
                    # this in then the call to getDrawRequests will think
                    # that we are still the instancer shape and not the 
                    # instance shape.  
                    # 
                    myinfo = OpenMayaUI.MDrawInfo( info )
                    myinfo.setMultiPath( geom.instanceShape )
                    shapeUI.getDrawRequests( myinfo,
                                             objectAndActiveOnly, myQueue )
                    geom.drawQueueList.append( (myQueue, position) )

        info.setMultiPath( path )
        # Finally we must supply a material back to the drawing code. 
        # We attempt to use the instance shape material; however, if 
        # that fails then we fall back to the default material 
        # 
        if shadedMode:
            defaultMaterial = OpenMayaUI.MMaterial.defaultMaterial()
            if not mainMaterial:
                mainMaterial = defaultMaterial 
            try: 
                request.setMaterial( mainMaterial )
            except:
                request.setMaterial( defaultMaterial ) 

            
        self.getDrawData(geom, data)            
        request.setDrawData(data)

        self.geometry = geom
        queue.add(request)

    # override
    def draw(self, request, view):
        """
         From the given draw request, get the draw data and determine
         which basic shape to draw and with what values.
        """
        
        data = request.drawData()
        shapeNode = self.surfaceShape()
        geom = self.geometry
        glFT.glMatrixMode( OpenMayaRender.MGL_MODELVIEW )
        if geom.instanceShape: 
            shapeUI = OpenMayaMPx.MPxSurfaceShapeUI.surfaceShapeUI(geom.instanceShape)
            for (queue, pos) in geom.drawQueueList:
                glFT.glPushMatrix();
                glFT.glTranslatef( pos[0], pos[1], pos[2] )
                while not queue.isEmpty():
                    request = queue.remove()
                    shapeUI.draw( request, view )
                glFT.glPopMatrix()

        # Draw a shell area that shows where the instances are being
        # drawn. This is nice to have if we don't have any instance
        # shapes connected to this plugin. 
        # 
        glFT.glPushAttrib( OpenMayaRender.MGL_ALL_ATTRIB_BITS )
        glFT.glPolygonMode(OpenMayaRender.MGL_FRONT_AND_BACK,
                           OpenMayaRender.MGL_LINE)
        glFT.glBegin(OpenMayaRender.MGL_QUADS)
        glFT.glVertex3f(-1*(geom.radius), -1*(geom.radius), 0.0)
        glFT.glNormal3f(0, 0, 1.0)
            
        glFT.glVertex3f(-1*(geom.radius), (geom.radius), 0.0)
        glFT.glNormal3f(0, 0, 1.0)
            
        glFT.glVertex3f((geom.radius), (geom.radius), 0.0)
        glFT.glNormal3f(0, 0, 1.0)
            
        glFT.glVertex3f((geom.radius), -1*(geom.radius), 0.0)
        glFT.glNormal3f(0, 0, 1.0)
        glFT.glEnd()
        glFT.glPopAttrib( )

    # override
    def select(self, selectInfo, selectionList, worldSpaceSelectPts):
        """
         Select function. Gets called when the bbox for the object is
         selected.  This function just selects the object without
         doing any intersection tests.
        """
        
        priorityMask = OpenMaya.MSelectionMask(OpenMaya.MSelectionMask.kSelectObjectsMask)
        item = OpenMaya.MSelectionList()
        item.add(selectInfo.selectPath())
        xformedPt = OpenMaya.MPoint()
        selectInfo.addSelection(item, xformedPt, selectionList,
                                 worldSpaceSelectPts, priorityMask, False)
        return True




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

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


def uiCreator():
    return OpenMayaMPx.asMPxPtr( instanceShapeUI() )


def nodeInitializer():
    # utility func for numeric attrs
    def setOptions(attr):
        attr.setHidden(False)
        attr.setKeyable(True)
        attr.setInternal(True)

    messageAttr = OpenMaya.MFnMessageAttribute() 
    numericAttr = OpenMaya.MFnNumericAttribute()

    instanceShape.aInstanceShape = messageAttr.create("instanceShape", "is")
    instanceShape.addAttribute(instanceShape.aInstanceShape)
    
    instanceShape.aRadius = numericAttr.create("radius", "r", OpenMaya.MFnNumericData.kDouble, kDefaultRadius)
    setOptions(numericAttr)
    instanceShape.addAttribute(instanceShape.aRadius)

    instanceShape.aCount = numericAttr.create("count", "ct", OpenMaya.MFnNumericData.kInt, kDefaultCount)
    setOptions(numericAttr)
    instanceShape.addAttribute(instanceShape.aCount)
    
# initialize the script plug-in
def initializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject, "Autodesk", "2011", "Any")
    try:
        mplugin.registerShape( kPluginNodeTypeName, spInstanceShapeNodeId,
                                nodeCreator, nodeInitializer, uiCreator )
    except:
        sys.stderr.write( "Failed to register node: %s" % kPluginNodeTypeName )
        raise


# uninitialize the script plug-in
def uninitializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject)
    try:
        mplugin.deregisterNode( spInstanceShapeNodeId )
    except:
        sys.stderr.write( "Failed to deregister node: %s" % kPluginNodeTypeName )
        raise