Detecting wire color in Maya III

This is te final post regarding color detection as described more broadly in those posts:
Wirecolor Part 1
Wirecolor Part 2
Detecting ffection

What I wish to leave you with is a documented piece of code which you can use. To me this is Mutils.color.py

'''
Created on Feb 18, 2013
@author: Trevor van Hoof
@package: Mutils
'''


from maya import cmds
import maya.utils


def affectsAll( attr, type ):
    '''
    Maps the affects net within a node,
    instead of just using cmds.affects
    this iterates over the resulting attributes
    again until all attributes indirectly
    affecting the given inAttr are found.
    
    @param inAttr: str, name of the attribute
    to map affecting attributes for.
    
    @param inType: str, type of the node we're
    looking at, this nodetype must of course
    have the inAttr.
    
    @returns: list of attribute names
    that affect the given inAttr.
    '''
    #these lists can in theory be precalculated constants
    attrs = cmds.affects(attr.rsplit('.',1)[-1], t=type)
    if not attrs:
        return []
    i = 0
    while i < len(attrs):
        tmp = cmds.affects(attrs[i], t=type) 
        if tmp:
            attrs.extend(tmp)
        attrs = list(set(attrs))
        i += 1
    return attrs


def affectedNet( inAttr, inNode ):
    '''
    This iteratively maps the affected network of the given
    attribute on the given node. The type of the node is 
    important and the attribute must exist on the node.
    
    It works by finding which internal inputs affect the given
    attribute, then it lists all connections to these attributes.
    
    For all node.attrs again the affected network is mapped
    until we have the entire node graph plus names of all attributes
    affecting the given inNode.inAttr through a connection in the DG.
    
    @returns: tuple of 2 lists, first containing the node names, second
    containing lists with attribute names (inputs and outputs) on that
    node that affect inNode.inAttr or connections thereto.
    
    The two lists have matching indices so the list of attribute names
    in returnValue[1][i] maps to the node name in returnValue[0][i]
    
    @param inAttr: str, name of the attribute to find affected net for
    @param inNode: str, path of the node to list inputs from 
    '''
    nodes = [inNode]
    attributes = [[inAttr]]
    
    #iterate until affection found or entire network traversed
    i = 0
    while i < len(nodes):
        #find internel affection net
        attributes[i].extend( affectsAll(attributes[i][0], cmds.nodeType(nodes[i])) )
        
        #find nodes that are connected to plugs in the affected net
        inputs = cmds.listConnections(nodes[i], s=True, d=False, c=True, p=True)
        if inputs:
            for j in range(0,len(inputs),2):
                #attribute name in affectednet
                if inputs[j].rsplit('.',1)[-1] in attributes[i]:
                    #get node attribute pair
                    nodeattr = inputs[j+1].split('.',1)
                    nodeattr[0] = cmds.ls(nodeattr[0], l=True)[0]
                    if nodeattr[0] not in nodes:
                        #append new nodes
                        nodes.append(nodeattr[0])
                        attributes.append([nodeattr[1]])
                    else:
                        #append new plugs on known nodes
                        attributes[ nodes.index(nodeattr[0]) ].append( nodeattr[1] )
        
        #if no incoming node was selected, continue iterating
        i += 1
    return nodes, attributes


def isAffected(inPathStr):
    '''
    This function grabs the affected network of the given node's matrix
    or output shape and checks whether any attributes of this network are driven
    by a selected node, or child of a selected node.
    
    It currently only supports geometry shapes.
    
    @param inPathStr: str, the node to check for
    
    @returns: bool, True if the given node is affected by one of the selected nodes
    '''
    #assume node is a transform by default
    attrib = 'matrix'
    
    #get the output attribute if node is a shape
    if cmds.ls(inPathStr, type='shape'):
        #detect the attribute name to get the affectedNet for
        nodetype = cmds.nodeType( inPathStr )
        if nodetype == 'mesh':
            attrib = 'outMesh'
        elif nodetype == 'subdiv':
            attrib = 'outSubdiv'
        elif nodetype in ('nurbsCurve','nurbsSurface'):
            attrib = 'local'
        else:
            raise ValueError('Nodetype %s of node %s not supported in isAffected'%(nodetype, inPathStr))
    elif not cmds.ls(inPathStr, type='dagNode'):
        raise ValueError('Given node path %s is not a Dag node in isAffected'%inPathStr)

    
    for node in affectedNet(attrib, inPathStr)[0]:
        if isParentSelected(node):
            return True
    return False


def isAffectedRecursively(inPathStr):
    '''
    Maps the affected net and checks if
    nodes in it are selected, if not,
    repeats the process for parents of the
    given node
    
    @param inPathStr: str, the node to check for
    
    @returns: bool, True if node or parent node
    is affected by a selected object
    '''
    obj = cmds.ls(inPathStr, l=True)
    if not obj:
        return False
    obj = obj[0]
    while obj and len(obj) > 1:
        if isAffected(obj):
            return True
        obj = obj.rsplit('|',1)[0]
    return False
    

def displayColorType(inObj):
    '''
    Returns the display color type, used by the
    cmds.displayColor() function, of the given node
    @todo: finish parsing imaginable node types
    
    @param inObj: node to get display color type name for
    
    @returns str: node type
    '''
    objtype = cmds.nodeType(inObj)
    if objtype == 'nurbsSurface':
        trims = cmds.listConnections(shape, s=True, d=False, type='planarTrimSurface')
        if trims:
            obtype = 'trimmedSurface'
        else:
            objtype = 'surface'
    if objtype == 'nurbsCurve':
        projectCurves = cmds.listConnections(shape, s=True, d=False, type='projectCurve')
        if projectCurves:
            objtype = 'curveOnSurface'
        else:
            objtype = 'curve'
    if objtype == 'mesh':
        objtype = 'polymesh'
    if objtype == 'joint' and cmds.listRelatives(shape, ad=True, type='effector'):
        objtype = 'segment'
    if objtype == 'cluster':
        objtype = 'locator'
    if objtype == 'distanceDimShape':
        objtype = 'dimension'
    return objtype
    
    
def isParentSelected(inObj, ignoreSelf=False):
    '''
    @param inPathStr: str, node to check parents for
    
    @param ignoreSelf: when set to True only the parents
    are checked for selection and not the input node
    
    @returns: bool, True when the given node or
    any of it's parents is selected
    '''
    selection = cmds.ls(sl=True, l=True)
    if not selection: #no selection, no result
        return
    if not ignoreSelf:
        if inObj in selection:
            return inObj
    targets = cmds.listRelatives(inObj, ap=True, f=True)
    if not targets:
        return
    for target in targets:
        if target in selection:
            return target
    return


def overrideAttr(inObj, inAttr):
    '''
    Gets the value of the given override attribute,
    searches in parents if overrides on the given object
    are not enabled, returns None if no overrides found
     
    @param inObj: str, node to start lookin from
    
    @param inAttr: str, attribute to find override value for
    
    @returns: value of the (parents) attribute or None
    '''
    target = inObj
    while target:
        if not cmds.getAttr('%s.overrideEnabled'%target):
            target = cmds.listRelatives(target, p=True, f=True)[0]
        return cmds.getAttr('%s.%s'%(target,inAttr))

    
def drawColor(inObj):
    '''
    Gets the color of the object's type. If the given object
    is not a shape node it returns the color of the first
    valid shape node directly below the given node
    
    @param inObjStr: string representing the path to the node
    to search for. If multiple nodes with the given name/path
    exist only the first will be used
        
    @returns: float[3], list of 0 to 1 RGB values
    '''
    #using executeInMainThreadWithResult to resolve 'bool is not a bool' errors
    #that should only occur when threading but still occur randomly all the time
    shapes = maya.utils.executeInMainThreadWithResult( 'cmds.listRelatives(\'%s\', ad=True, type=\'shape\', f=True)'%inObj )
    if not shapes:
        if cmds.nodeType(inObj) != 'transform':
            shape = inObj
        else: #transform node without shapes has no color
            return None
    else:
        shape = shapes[0]

    nodetype = displayColorType( shape )
    selected = isParentSelected( shape )
    displaytype = overrideAttr(shape, 'overrideDisplayType')
    
    if selected:
        #templated
        if displaytype == 1:
            return cmds.colorIndex( cmds.displayColor('activeTemplate', q=True, active=True), q=True )
        #lead
        if selected == cmds.ls(os=True, l=True)[-1]:
            return cmds.colorIndex( cmds.displayColor('lead', q=True, active=True), q=True )
            
        #active
        return cmds.colorIndex( cmds.displayColor(nodetype, q=True, active=True), q=True )
        
    #affected
    if cmds.displayPref( q=True, displayAffected=True ) and isAffectedRecursively( shape ):
        #if obj is affected by something that is selected
        return cmds.colorIndex( cmds.displayColor('activeAffected', q=True, active=True), q=True )
    
    #referenced
    if displaytype == 2:
        return cmds.colorIndex( cmds.displayColor('referenceLayer', q=True), q=True )
        
    #templated
    if displaytype == 1:
        return cmds.displayRGBColor('template', q=True)
    
    #override color
    overridecolor = overrideAttr(shape, 'overrideColor')
    if overridecolor: #not None and not 0
        return cmds.colorIndex( overridecolor, q=True )

    #dormant
    return cmds.colorIndex( cmds.displayColor(nodetype, q=True, dormant=True), q=True )

And a little demonstration, which requires PyQt4. It draws the current wire color of the current target. Push the button to set selected object as new target ( cmds.ls(sl=True)[0] ), updates color when you change selection, so it may be a tad confusing but you have to keep selecting stuff to see it change; which mostyly works but not when adding constraints, changing layer settings or changing target, then just deselect and undo or use the up and down arrow keys or something.

from maya import cmds
from Mutils import color
reload(color)
from PyQt4 import QtCore, QtGui
import sip
from maya import OpenMayaUI
import maya.utils

mainwindow = sip.wrapinstance( long(OpenMayaUI.MQtUtil.mainWindow()), QtGui.QMainWindow )

class coloredRect( QtGui.QDockWidget ):
    def __init__(self):
        QtGui.QFrame.__init__(self, mainwindow)
        
        btn = QtGui.QPushButton("Show selected",self)
        btn.clicked.connect(self.storeobj)
        self.storeobj(None)
        
        self._job = cmds.scriptJob(e=['SelectionChanged', self.updatebrush])
        self._brush = QtCore.Qt.NoBrush
        
        self.updatebrush()
        
        self.setFloating(True)
        self.show()
        
    def storeobj(self, e):
        self._obj = maya.utils.executeInMainThreadWithResult( 'cmds.ls(sl=True, l=True)' )
        print self._obj
        
    def paintEvent(self, e):
        if self._brush != QtCore.Qt.NoBrush:
            r = self.geometry()
            r.setTop(0)
            r.setLeft(0)
            painter = QtGui.QPainter(self)
            painter.setBrush(self._brush)
            painter.drawRect(r)

    def updatebrush(self):
        if self._obj:
            c = color.drawColor(self._obj[0])
            print c
            c = QtGui.QColor( c[0]*255, c[1]*255, c[2]*255, 255 )
            self._brush = QtGui.QBrush( c )
        else:
            self._brush = QtCore.Qt.NoBrush
        self.repaint()

    def __del__(self):
        cmds.scriptJob(k=self._job, force=True)

try:
    if not cmds.objExists( c ):
        raise
except:
    c = cmds.polyCube()[0]
    s = cmds.polySphere()[0]
    cmds.xform(s,t=[0,0,2])
    cn = cmds.orientConstraint(c, s)
finally:
    cmds.scriptJob(ka=True)
    w = coloredRect()

    cmds.select(c)

color.drawColor(w._obj[0])
w.update()
w.repaint()

Leave a Reply

Your email address will not be published. Required fields are marked *