mrv.maya.nt.it

1 # -*- coding: utf-8 -*- 2 """ 3 Contains different multi-purpose iterators allowing to conveniently walk the dg and 4 dag. 5 """ 6 __docformat__ = "restructuredtext" 7 8 import maya.OpenMaya as api 9 import maya.cmds as cmds 10 from maya.OpenMaya import MDagPath, MObject 11 from base import Node, DagNode, NodeFromObj, Component 12 13 __all__ = ("dgIterator", "dagIterator", "graphIterator", "selectionListIterator", 14 "iterDgNodes", "iterDagNodes", "iterGraph", "iterSelectionList") 15

16 -def _argsToFilter( args ):

17 """convert the MFnTypes in args list to the respective typeFilter""" 18 typeFilter = api.MIteratorType( ) 19 if args: 20 if len(args) == 1 : 21 typeFilter.setFilterType ( args[0] ) 22 else : 23 # annoying argument conversion for Maya API non standard C types 24 scriptUtil = api.MScriptUtil() 25 typeIntM = api.MIntArray() 26 scriptUtil.createIntArrayFromList( args, typeIntM ) 27 typeFilter.setFilterList( typeIntM ) 28 # we will iterate on dependancy nodes, not dagPaths or plugs 29 typeFilter.setObjectType( api.MIteratorType.kMObject ) 30 # create iterator with (possibly empty) typeFilter 31 return typeFilter

32 33 34 #{ Iterator Creators 35

36 -def dgIterator( *args, **kwargs ):

37 """ 38 :return: MItDependencyNodes configured according to args - see docs at 39 `iterDgNodes`. 40 :note: use this method if you want to use more advanced features of the iterator""" 41 typeFilter = _argsToFilter( args ) 42 iterObj = api.MItDependencyNodes( typeFilter ) 43 return iterObj

44

45 -def dagIterator( *args, **kwargs ):

46 """ 47 :return: MItDagIterator configured according to args - see docs at 48 `iterDagNodes`. 49 :note: use this method if you want to use more advanced features of the iterator""" 50 depth = kwargs.get('depth', True) 51 underworld = kwargs.get('underworld', False) 52 root = kwargs.get('root', None ) 53 typeFilter = _argsToFilter( args ) 54 55 # SETUP TYPE FILTER - reset needs to work with root 56 if root is not None: 57 if isinstance( root, (MDagPath, DagNode) ): 58 typeFilter.setObjectType( api.MIteratorType.kMDagPathObject ) 59 else : 60 typeFilter.setObjectType( api.MIteratorType.kMObject ) 61 62 # create iterator with (possibly empty) filter list and flags 63 if depth : 64 traversal = api.MItDag.kDepthFirst 65 else : 66 traversal = api.MItDag.kBreadthFirst 67 68 iterObj = api.MItDag( typeFilter, traversal ) 69 70 # set start object 71 if root is not None : 72 startObj = startPath = None 73 if isinstance( root, MDagPath): 74 startPath = root 75 elif isinstance( root, DagNode ): 76 startPath = root.dagPath() 77 elif isinstance( root, Node ): 78 startObj = root.object() 79 else: 80 startObj = root 81 # END handle obj type 82 iterObj.reset( typeFilter, startObj, startPath, traversal ) 83 # END if root is set 84 85 if underworld : 86 iterObj.traverseUnderWorld( True ) 87 else : 88 iterObj.traverseUnderWorld( False ) 89 90 91 return iterObj

92 93

94 -def graphIterator( nodeOrPlug, *args, **kwargs ):

95 """ 96 :return: MItDependencyGraph configured according to args - see docs at 97 `iterGraph`. 98 :note: use this method if you want to use more advanced features of the iterator 99 :raise RuntimeError: if the filter types does not allow any nodes to be returned. 100 This is a bug in that sense as it should just return nothing. It also shows that 101 maya pre-parses the result and then just iterates over a list with the iterator in 102 question""" 103 startObj = startPlug = None 104 105 if isinstance( nodeOrPlug, api.MPlug ): 106 startPlug = nodeOrPlug 107 startObj = MObject() 108 elif isinstance( nodeOrPlug, Node ): 109 startObj = nodeOrPlug.object() 110 startPlug = nullplugarray[0] 111 else: 112 startObj = nodeOrPlug 113 startPlug = nullplugarray[0] 114 # END traversal root 115 116 inputPlugs = kwargs.get('input', False) 117 breadth = kwargs.get('breadth', False) 118 plug = kwargs.get('plug', False) 119 prune = kwargs.get('prune', False) 120 typeFilter = _argsToFilter( args ) 121 122 if startPlug is not None : 123 typeFilter.setObjectType( api.MIteratorType.kMPlugObject ) 124 else : 125 typeFilter.setObjectType( api.MIteratorType.kMObject ) 126 # END handle object type 127 128 direction = api.MItDependencyGraph.kDownstream 129 if inputPlugs : 130 direction = api.MItDependencyGraph.kUpstream 131 132 traversal = api.MItDependencyGraph.kDepthFirst 133 if breadth : 134 traversal = api.MItDependencyGraph.kBreadthFirst 135 136 level = api.MItDependencyGraph.kNodeLevel 137 if plug : 138 level = api.MItDependencyGraph.kPlugLevel 139 140 iterObj = api.MItDependencyGraph( startObj, startPlug, typeFilter, direction, traversal, level ) 141 142 iterObj.disablePruningOnFilter() 143 if prune : 144 iterObj.enablePruningOnFilter() 145 146 return iterObj

147 148

149 -def selectionListIterator( sellist, **kwargs ):

150 """ 151 :return: iterator suitable to iterate given selection list - for more info see 152 `iterSelectionList`""" 153 filtertype = kwargs.get( "filterType", api.MFn.kInvalid ) 154 iterator = api.MItSelectionList( sellist, filtertype ) 155 return iterator

156 157 #} END iterator creators 158 159

160 -def iterDgNodes( *args, **kwargs ):

161 """ Iterator on MObjects or Nodes of the specified api.MFn types 162 163 :param args: type as found in MFn.k... to optionally restrict the set of nodes the iterator operates upon. 164 All nodes of a type included in the args will be iterated on. 165 args is empty, all nodes of the scene will be iterated on which may include DAG nodes as well. 166 :param kwargs: 167 * asNode: 168 if True, default True, the returned value will be wrapped as node 169 * predicate: 170 returns True for every iteration element that may be returned by the iteration, 171 default : lambda x: True""" 172 iterator = dgIterator( *args, **kwargs ) 173 predicate = kwargs.get( "predicate", lambda x: True ) 174 asNode = kwargs.get( "asNode", True ) 175 176 isDone = iterator.isDone 177 thisNode = iterator.thisNode 178 next = iterator.next 179 180 while not isDone() : 181 node = thisNode() 182 if asNode: 183 node = NodeFromObj( node ) 184 if predicate( node ): 185 yield node 186 next()

187 # END for each obj in iteration 188 189 # Iterators on dag nodes hierarchies using MItDag (ie listRelatives)

190 -def iterDagNodes( *args, **kwargs ):

191 """ Iterate over the hierarchy under a root dag node, if root is None, will iterate on whole Maya scene 192 If a list of types is provided, then only nodes of these types will be returned, 193 if no type is provided all dag nodes under the root will be iterated on. 194 Types are specified as Maya API types being a member of api.MFn 195 The following keywords will affect order and behavior of traversal: 196 197 :param kwargs: 198 * dagpath: 199 if True, default True, MDagPaths will be returned 200 If False, MObjects will be returned - it will return each object only once in case they 201 occour in multiple paths. 202 * depth: 203 if True, default True, Nodes will be returned as a depth first traversal of the hierarchy tree 204 if False as a post-order (breadth first) 205 * underworld: 206 if True, default False, traversal will include a shape's underworld 207 (dag object parented to the shape), if False the underworld will not be traversed, 208 * asNode: 209 if True, default True, the returned item will be wrapped into a Node 210 * root: 211 MObject or MDagPath or Node of the object you would like to start iteration on, or None to 212 start on the scene root. The root node will also be returned by the iteration ! 213 Please note that if an MObject is given, it needs to be an instanced DAG node to have an effect. 214 * predicate: 215 method returning True if passed in iteration element can be yielded 216 default: lambda x: True""" 217 218 # Must define dPath in loop or the iterator will yield 219 # them as several references to the same object (thus with the same value each time) 220 # instances must not be returned multiple times 221 # could use a dict but it requires "obj1 is obj2" and not only "obj1 == obj2" to return true to 222 iterator = dagIterator( *args, **kwargs ) 223 isDone = iterator.isDone 224 next = iterator.next 225 226 dagpath = kwargs.get('dagpath', True) 227 asNode = kwargs.get('asNode', True ) 228 predicate = kwargs.get('predicate', lambda x: True ) 229 230 if dagpath: 231 getPath = iterator.getPath 232 while not isDone( ) : 233 rval = MDagPath( ) 234 getPath( rval ) 235 if asNode: 236 rval = NodeFromObj( rval ) 237 if predicate( rval ): 238 yield rval 239 240 next() 241 # END while not is done 242 # END if using dag paths 243 else: 244 # NOTE: sets don't work here, as more than == comparison is required 245 instanceset = list() 246 currentItem = iterator.currentItem 247 isInstanced = iterator.isInstanced 248 249 while not isDone() : 250 rval = currentItem() 251 if isInstanced( True ): 252 if rval not in instanceset: 253 instanceset.append( rval ) 254 else: 255 next() 256 continue 257 # END if object not yet returned 258 # END handle instances 259 260 if asNode: 261 rval = NodeFromObj(rval) 262 if predicate( rval ): 263 yield rval 264 265 next()

266 # END while not is done 267 # END if using mobjects 268

269 -def iterGraph( nodeOrPlug, *args, **kwargs ):

270 """ Iterate Dependency Graph (DG) Nodes or Plugs starting at a specified root Node or Plug. 271 The iteration _includes_ the root node or plug. 272 The following keywords will affect order and behavior of traversal: 273 274 :param nodeOrPlug: Node, MObject or MPlug to start the iteration at 275 :param args: list of MFn node types 276 If a list of types is provided, only nodes of these types will be returned, 277 if no type is provided all connected nodes will be iterated on. 278 :param kwargs: 279 * input: 280 if True connections will be followed from destination to source, 281 if False from source to destination 282 default is False (downstream) 283 * breadth: 284 if True nodes will be returned as a breadth first traversal of the connection graph, 285 if False as a preorder (depth first) 286 default is False (depth first) 287 * plug: 288 if True traversal will be at plug level (no plug will be traversed more than once), 289 if False at node level (no node will be traversed more than once), 290 default is False (node level) 291 * prune: 292 if True, the iteration will stop on nodes that do not fit the types list, 293 if False these nodes will be traversed but not returned 294 default is False (do not prune) 295 * asNode: 296 if True, default True, and if the iteration is on node level, 297 Nodes ( wrapped MObjects ) will be returned 298 If False, MObjects will be returned 299 * predicate: 300 method returning True if passed in iteration element can be yielded 301 default: lambda x: True 302 :return: Iterator yielding MObject, Node or Plug depending on the configuration flags, first yielded item is 303 always the root node or plug.""" 304 try: 305 iterator = graphIterator( nodeOrPlug, *args, **kwargs ) 306 except RuntimeError: 307 # may raise if iteration would yield no results 308 raise StopIteration() 309 310 retrievePlugs = not iterator.atNodeLevel( ) 311 asNode = kwargs.get( "asNode", True ) 312 predicate = kwargs.get( 'predicate', lambda x: True ) 313 314 isDone = iterator.isDone 315 next = iterator.next 316 thisPlug = iterator.thisPlug 317 currentItem = iterator.currentItem 318 319 # iterates and yields MObjects 320 rval = None 321 # if node filters are used, it easily threw NULL Object returned errors 322 # just because the iteration is depleted - catching this now 323 try: 324 while not isDone(): 325 if retrievePlugs: 326 rval = thisPlug() 327 else: 328 rval = currentItem() 329 if asNode: 330 rval = NodeFromObj( rval ) 331 # END handle asNode 332 # END if return on node level 333 334 if predicate( rval ): 335 yield rval 336 337 next() 338 # END of iteration 339 except RuntimeError: 340 raise StopIteration()

341 # END handle possible iteration error 342 343 344 nullplugarray = api.MPlugArray() 345 nullplugarray.setLength( 1 )

346 -def iterSelectionList( sellist, filterType = api.MFn.kInvalid, predicate = lambda x: True, 347 asNode = True, handlePlugs = True, handleComponents = False ):

348 """Iterate the given selection list 349 350 :param sellist: MSelectionList to iterate 351 :param filterType: MFnType id acting as simple type filter to ignore all objects which do not 352 have the given object type 353 :param asNode: if True, returned MObjects or DagPaths will be wrapped as Node, compoents will be 354 wrapped as Component. 355 Otherwise they will be returned as MObjects and MDagPaths respectively. 356 :param handlePlugs: if True, plugs can be part of the selection list and will be returned. This 357 implicitly means that the selection list will be iterated without an iterator, and MFnType filters 358 will be slower as it is implemented in python. If components are enabled, the tuple returned will be 359 ( Plug, MObject() ) 360 :param predicate: method returninng True if passed in iteration element can be yielded 361 default: lambda x: True 362 :param handleComponents: if True, possibly selected components of dagNodes will be returned 363 as well. This forces the return value into tuple(Node, Component) 364 :return: Node or Plug on each iteration step 365 If handleComponents is True, for each Object, a tuple will be returned as tuple( Node, Component ) where 366 component is NullObject ( MObject ) if the whole object is on the list. 367 If the original object was a plug, it will be in the tuples first slot, whereas the component 368 will be a NullObject""" 369 kNullObj = MObject() 370 if handlePlugs: 371 # version compatibility - maya 8.5 still defines a plug ptr class that maya 2005 lacks 372 plug_types = api.MPlug 373 if cmds.about( v=1 ).startswith( "8.5" ): 374 plug_types = ( api.MPlug, api.MPlugPtr ) 375 376 # SELECTION LIST MODE 377 kInvalid = api.MFn.kInvalid 378 getDagPath = sellist.getDagPath 379 getPlug = sellist.getPlug 380 getDependNode = sellist.getDependNode 381 for i in xrange( sellist.length() ): 382 # DAG PATH 383 rval = None 384 component = kNullObj 385 try: 386 rval = MDagPath( ) 387 if handleComponents: 388 component = MObject() 389 getDagPath( i, rval, component ) 390 if asNode and not component.isNull(): 391 component = Component( component ) 392 # END handle asNode 393 else: 394 getDagPath( i, rval ) 395 # END handle components in DagPaths 396 except RuntimeError: 397 # TRY PLUG - first as the object could be returned as well if called 398 # for DependNode 399 try: 400 rval = nullplugarray[0] 401 getPlug( i, rval ) 402 # try to access the attribute - if it is not really a plug, it will 403 # fail and throw - for some reason maya can put just the depend node into 404 # a plug 405 rval.attribute() 406 except RuntimeError: 407 # TRY DG NODE 408 rval = MObject( ) 409 getDependNode( i, rval ) 410 # END its not an MObject 411 # END handle dagnodes/plugs/dg nodes 412 413 # should have rval now 414 if isinstance( rval, plug_types ): 415 # apply filter 416 if filterType != kInvalid and rval.node().apiType() != filterType: 417 continue 418 # END apply filter type 419 else: 420 if filterType != kInvalid: 421 # must be MDagPath or MObject 422 if rval.apiType() != filterType: 423 continue 424 # END filter handling 425 426 if asNode: 427 rval = NodeFromObj( rval ) 428 # END plug handling 429 430 if handleComponents: 431 rval = ( rval, component ) 432 433 if predicate( rval ): 434 yield rval 435 # END for each element 436 else: 437 # ITERATOR MODE 438 # the code above can handle it all, this one might be faster though 439 iterator = selectionListIterator( sellist, filterType = filterType ) 440 kDagSelectionItem = api.MItSelectionList.kDagSelectionItem 441 kDNselectionItem = api.MItSelectionList.kDNselectionItem 442 rval = None 443 444 isDone = iterator.isDone 445 itemType = iterator.itemType 446 getDagPath = iterator.getDagPath 447 getDependNode = iterator.getDependNode 448 next = iterator.next 449 while not isDone(): 450 # try dag object 451 component = kNullObj 452 itemtype = itemType( ) 453 if itemtype == kDagSelectionItem: 454 rval = MDagPath( ) 455 if handleComponents: 456 component = MObject( ) 457 getDagPath( rval, component ) 458 if asNode and not component.isNull(): 459 component = Component( component ) 460 # END handle component conversion 461 else: 462 getDagPath( rval ) 463 # END handle components 464 else: 465 rval = MObject() 466 getDependNode( rval ) 467 # END handle item type 468 469 if asNode: 470 rval = NodeFromObj( rval ) 471 # END handle as node 472 473 if handleComponents: 474 rval = ( rval, component ) 475 # END handle component 476 477 if predicate( rval ): 478 yield rval 479 480 next()

481 # END while not done 482

Source Code for Module mrv.maya.nt.it