Renderable Skeleton - Maya 2015 PyMel

'''
    Script by Dr. Gravitas

    This script takes one or more joint selections and creates a series of spheres and cones
    along the joint heirarchies to represent the skeletons with renderable components.
'''
import pymel.core as pm
from maya import OpenMaya as om

'''
    For a given joint, create a sphere and cone(s) representing the joint and bone lengths to any child joints
'''
def placeBone( aJoint, shaderGroup ):
    # We can't use | in names, as would happen if two joints share the same name,
    #   but just using the .nodeName() name will not be unique. So replace | with empty String
    jointName = aJoint.name().replace('|', '')
    #create a PolySphere with a name derived from the joint...
    sphereName = 'VisSkel_' + jointName + '_sphere'

    # (Joint Radius is a bit big; scale down by 0.09. It'd be nice if we could use user joint size display prefs to calc this.)
    sphere = pm.nodetypes.PolySphere( n=sphereName, sx=6, sy=6, r=aJoint.getRadius()*0.09 )
    # ...get the sphere's transform node through selection...
    #       (The name given above is applied to the transform node, rather than it's PolySphere node because Maya!)
    sphTrans = getTransformNodeFromPolyPrimative( sphereName )
    # ...and parent it the joint, inheriting the rotation/Position, then set the rest of joint's transformation attributes.
    pm.parent( sphTrans, aJoint, relative=True )
    sphTrans.setScale( aJoint.getScale() )

    # Add sphere to the given shader group; eg. applying the shader to the sphere.
    #shaderGroup.addMember( sphTrans )
    pm.sets( shaderGroup, edit = True, forceElement=sphTrans )
    #The hard part: creating cones pointing to any child joints of the current joint.
    placeCones( aJoint, jointName, shaderGroup )

'''
    Get a transformNode of a PolyPrimative that was created with a name (ex. PolyCone, PolySphere, etc.)
    This method will clear your selections.
    TODO: Figure out a better, to do this. Using select is a clunky work around
'''
def getTransformNodeFromPolyPrimative( polyPrimName ):
    pm.select( polyPrimName )
    transformNode = (pm.selected())[0]
    pm.select(cl=True)
    return transformNode

'''
    Places Cones representing the bone length between a given joint and its children.
'''
def placeCones( joint, jointName, shaderGroup ):
    for child in joint.getChildren(type="joint"):
        # create the cone;
        coneName = 'VisSkel_' + child.name().replace('|', '') + '_cone_' + jointName
        # cone length is the distance between joint and it's child
        height = om.MVector(*child.attr('t').get()).length()
        cone = pm.polyCone(n=coneName, sh=1, sa=3, sc=0, r=joint.getRadius()*0.09,height=height )

        # position and orient the cone
        pm.delete(pm.pointConstraint(joint, cone))
        pm.delete(pm.aimConstraint(child, cone, aim=[0,1,0]))
        pm.move(0, height/2, 0, cone, objectSpace=True, r=True)

        coneTrans = getTransformNodeFromPolyPrimative( coneName )
        pm.parent( coneTrans, joint )

        #shaderGroup.addMember( coneTrans )
        pm.sets( shaderGroup, edit = True, forceElement=coneTrans )

'''
    Create A Lambert Shader of the given color, with the name VisSkel_Lambert# where # is
    the number in the order of such nodes created.
'''
def createLambertShader( color ):
    # Create a Lambert shader of a given name and associated shaderGroup
    shader, shaderGroup = pm.createSurfaceShader( 'lambert', name='VisSkel_Lambert' );
    # Set the shader's color
    shader.setColor( color )
    return shader, shaderGroup

'''
    Create a unique identifier string representing a color for a dictionary key.
'''
def colorKey( color ):
    return str(color.r)+str(color.g)+str(color.b)+str(color.a)

'''
    This method takes a list of joint hierarchies and creates a series of spheres and cones
    to create a renderable skeleton. The method also creates shaders to apply to the visible
    skeleton, colored based on the colorized joint skeleton. If the skeleton was never bound,
    then all the colors will likely be blue. This method takes care to not reoperate on joints
    and will not duplicate shaders of the same color.
'''
def generateSkeleton( targets ):
    # Keep a dictionary of targets (and their child joints) that we've hit so we don't operate on them again.
    jointsHit = {}
    # Dictionary mapping Colors to their shader/shaderGroups so we do not duplicate shaders with the same color.
    shaderColors = {}

    for target in targets:
        # Select the target joint and all its children via selecting the target joint with heirarcy flag
        pm.select( target.name(), hi=True )
        # Get those selected joints then clear selection.
        jointHeirarchy = pm.selected()
        pm.select(cl=True)

        for aJoint in jointHeirarchy:
            # if this is not actually a joint, or if we've already operated on a joint with this name,
            #   then we just skip it without throwing error.
            if isinstance(aJoint, pm.nodetypes.Joint) and aJoint.name() not in jointsHit:
                # record joint as hit.
                jointsHit[aJoint.name()] = aJoint
                # Get the joint viewport display color so we have a colorized skeleton
                rawColor = aJoint.__apimfn__().dormantColor()
                color = pm.datatypes.Color( rawColor.r, rawColor.g, rawColor.b)

                key = colorKey( color )
                if key not in shaderColors:
                    # If the color does not exist as a shader yet, create one.
                    shader, shaderGroup = createLambertShader( color )
                    shaderColors[key] = shaderGroup
                else:
                    # Otherwise, simply get the existing shader
                    shaderGroup = shaderColors[key]

                placeBone( aJoint, shaderGroup )

# End of method definitions; Beginning of the script execution!
# Jump the command log output a couple lines to make clear the difference between code reprint and output.
print '\n\n\n'

#Get list of targets from what has been selected
targets = pm.selected()

# Verify all selected target are joints.
for target in targets:
    if not isinstance(target, pm.nodetypes.Joint):
        pm.error(target + ' is Type: ' + type(target).__name__ + '. Not the required Type: Joint')

# clear selections
pm.select(cl=True)

generateSkeleton( targets )