Maya: Matte Tool

matte_tool.py

import sys

import maya.cmds as cmds
import pymel
import pymel.core
from PySide import QtGui


class Example(QtGui.QWidget):
    
    def __init__(self):
        super(Example, self).__init__()

    def show(self):
        text, ok = QtGui.QInputDialog.getText(self, 'Input Dialog', 
            'Enter your name:')
        
        if ok:
            return str(text)

cl = pymel.core.nodetypes.RenderLayer.currentLayer()
sel = pymel.core.ls(selection=True)

def createMatte(nodes, name=None):
    rp = None
    if name:
        rp = pymel.core.general.createNode('renderPass')
    else:
        if pymel.core.objExists('matte_rp'):
            rp = pymel.core.PyNode('matte_rp')
        else:
            rp = pymel.core.general.createNode('renderPass')
    pcm = pymel.core.general.createNode('passContributionMap')
    
    print rp
    if name:
        pymel.core.rename(rp, name + '_rp')
        pymel.core.rename(pcm, name + '_pcm')
    else:
        pymel.core.rename(rp, 'matte_rp')
    
    cmds.setRenderPassType(rp.name(), type='MATTE')
    rp.frameBufferType.set(256)
    rp.filtering.set(0)
    
    cl.passContributionMap.connect(pcm.owner, na=True)
    cl.renderPass.connect(rp.owner, na=True)
    
    rp.message.connect(pcm.renderPass, na=True)
    
    for node in nodes:
        if node.nodeType() != 'transform':
            node.getParent().message.connect(pcm.dagObjects, na=True)
        else:
            node.message.connect(pcm.dagObjects, na=True)

for node in sel:
    if node.nodeType() == 'lambert':
        se = node.listConnections(type='shadingEngine')[0]
        transforms = se.listConnections(type='mesh')
        print node.name().split(':')[-1].split('_')[0]

reply = cmds.confirmDialog( title='Matte', message='One or more mattes?', button=['One','More'])

if reply == 'One':
    ex = Example()
    name = ex.show()
    
    if name:
        createMatte(sel, name)
    else:
        createMatte(sel)
else:
    for node in sel:
        
        if node.nodeType() == 'lambert':
            se = node.listConnections(type='shadingEngine')[0]
            transforms = se.listConnections(type='mesh')
            name = node.name().split(':')[-1].split('_')[0]
            
            createMatte(transforms, name)
        else:
            createMatte([node], node.name().split(':')[-1])