[Maya] Divide selected joint

divideJoint.py

import pymel.core as pm

# Helper methods
def debug(message):
    print("DEBUG: %s" % message)
    

# Average between 2 nodes
def averageBetween(a, b):
    return matrixLinearBetween(a, b, 0.5)

def matrixLinearBetween(a, b, fac):
    matA = pm.datatypes.Matrix(pm.xform(a, q=True, m=True, ws=True))
    matB = pm.datatypes.Matrix(pm.xform(b, q=True, m=True, ws=True))
    return linearBetween(matA, matB, fac)

def linearBetween(a, b, fac):
    return (a * (1 - fac)) + (b * fac)


def subdivideJoint(a, b, amount = 1):
    if amount < 1:
        return # Won't work

    # Connect to last created node
    lastNode = a

    # Stepsize
    steps = 1.0 / (amount + 1)

    # Loop through every division
    for i in range(amount):
        # Deselect everything
        pm.select(clear=True)

        # Brand new joint
        joint = pm.joint()

        # Get the linear factor of range 0 - 1
        factor = steps * (i + 1)

        # Set node between 2 other nodes
        pm.xform(joint, ws=True, m=matrixLinearBetween(firstNode, secondNode, factor))

        # Set aim to be from firstNode
        pm.xform(joint, ws=True, rotation= pm.xform(firstNode, q=True, rotation=True, ws=True) )

        # update radius
        averageRadius = linearBetween(firstNode.radius.get(), secondNode.radius.get(), factor)
        joint.radius.set(averageRadius)

        # Freeze rotation
        pm.makeIdentity( joint, apply=True, rotate=True ) 

        # Parent
        pm.parent(joint, lastNode)
        lastNode = joint
    pm.parent(secondNode, lastNode)



#
# Magic goes here
#
selection = pm.ls(sl=True, type="joint")
if len(selection) > 0:
    result = pm.promptDialog(
                    title='Subdivide Joint',
                    message='Amount of divisions:',
                    button=['OK', 'Cancel'],
                    defaultButton='OK',
                    cancelButton='Cancel',
                    dismissString='Cancel')

    if result == 'OK':
        divisions = int(pm.promptDialog(q=True, text=True))
        print divisions

        if len(selection) == 1:
            # temp nodes
            firstNode = selection[0]
            theKids = firstNode.getChildren()
            if( len(theKids) == 1): 
                secondNode = theKids[0]
                subdivideJoint(firstNode, secondNode, divisions)
            else:
                pm.warning("Too many branches, select a second joint.")

        elif len(selection) == 2:
            firstNode = selection[0]
            secondNode = selection[1]

            if secondNode in firstNode.getChildren():
                subdivideJoint(firstNode, secondNode, divisions)
            else:
                pm.warning("%s is no direct child of %s." % (firstNode, secondNode))

else:
    pm.warning("No joints were selected.")