import re
import time
def create_mattes():
    rst_prefix = '^RST_'
    set_suffix = 'ObjectSet$'
    all_sets = cmds.ls(set=1)

    sets = []
    for s in all_sets:
        if re.search(rst_prefix, s):
            sets.append(s)

    for set_node in sets:
        node_basename = re.sub('{0}|{1}'.format(rst_prefix, set_suffix), '', set_node)
        aov_name = 'matte_{0}'.format(node_basename)


        t_switch = cmds.shadingNode('tripleShadingSwitch', au=1)
        cmds.setAttr('{0}.default'.format(t_switch), 0, 0, 0, typ='double3')

        shader_node = cmds.shadingNode('aiUtility', asShader=1)
        # Set this to flat.
        cmds.setAttr('{0}.shadeMode'.format(shader_node), 2)
        # Set color to white.  This is default, but we'll keep it here.
        cmds.setAttr('{0}.color'.format(shader_node), 1, 1, 1, type='double3')

        cmds.connectAttr('{0}.output'.format(t_switch), '{0}.color'.format(shader_node, f=1))


        for obj in cmds.listRelatives(set_node, pa=1):
            start = time.time()
            index = cmds.getAttr('{0}.input'.format(t_switch), s=1)
            cmds.connectAttr('{0}.instObjGroups[0]'.format(obj), '{0}.input[{1}].inShape'.format(t_switch, index), f=1)
            cmds.connectAttr('{0}.outColor'.format(shader_node), '{0}.input[{1}].inTriple'.format(t_switch, index), f=1)
            print 'obj', time.time() - start

        # AOV'S
        aov_size = cmds.getAttr('defaultArnoldRenderOptions.aovList', s=1)

        custom_aov = cmds.createNode('aiAOV', n=aov_name, skipSelect=True)

        cmds.connectAttr('{0}.message'.format(custom_aov), 'defaultArnoldRenderOptions.aovList[{0}]'.format(str(aov_size)), f=1)

        cmds.connectAttr('defaultArnoldDriver.message', '{0}.outputs[0].driver'.format(custom_aov), f=1)
        cmds.connectAttr('defaultArnoldFilter.message', '{0}.outputs[0].filter'.format(custom_aov), f=1)

        # connect to default shader
        cmds.connectAttr('{0}.outColor'.format(shader_node), '{0}.defaultValue'.format(custom_aov), f=1)

sphere = cmds.polySphere()
maya_objects_set = cmds.sets(sphere, n='RST_testingObjectSet')

tmp_start = time.time()
create_mattes()
print 'Total Time', time.time() - tmp_start