Package mrv :: Package maya :: Package nt :: Module set
[hide private]
[frames] | no frames]

Source Code for Module mrv.maya.nt.set

  1  # -*- coding: utf-8 -*- 
  2  """ Contains improved clases for set and partition editing  """ 
  3  __docformat__ = "restructuredtext" 
  4   
  5  import base as nt 
  6  import typ 
  7  import maya.OpenMaya as api 
  8  import maya.cmds as cmds 
  9  import it 
 10  import mrv.maya.undo as undo 
11 12 13 #{ Exceptions 14 -class ConstraintError( RuntimeError ):
15 """Thrown if a partition does not allow objects to be added, and the addition 16 was not forced, and failure was not ignored as well"""
17 #}
18 19 20 -class ObjectSet:
21 """ Extended and more convenient object set interface dealing with Nodes ( and 22 provides the original MFnSet interface as well 23 """ 24 __metaclass__ = typ.MetaClassCreatorNodes 25 kReplace, kAdd,kAddForce, kRemove = range( 4 ) 26 27 #{ Partition Handling 28
29 - def partitions( self ):
30 """:return: list of Nodes of partitions the entity is set is part of""" 31 return [ p.mwrappedNode() for p in self.partition.moutputs() ]
32 33 34 @undoable
35 - def setPartition( self, partition, mode = 0 ):
36 """Add, add exclusive or remove the given partition from our partition list 37 38 :param partition: Node, representing the partition, or a list of such 39 :param mode: 40 * 0 = replace 41 * 1 = add 42 * 2 = remove 43 :return: self for chained operations 44 :note: use the supplied enumeration to specify the mode""" 45 # convert to list 46 prts = partition 47 if isinstance( partition, Partition ): 48 prts = [ partition ] 49 50 if mode == self.kReplace: 51 self.setPartition( self.partitions( ), self.kRemove ) 52 mode = self.kAdd # now the partitions have to be added 53 # go ahead with the add 54 # END replace mode 55 56 if mode == self.kRemove: 57 for part in prts: 58 self.partition.mdisconnectNode( part ) 59 return self 60 # END remove mode 61 62 if mode == self.kAdd: 63 # only allow to be connected once 64 for part in prts: 65 self.partition.mconnectToArray( part.st, exclusive_connection = True ) 66 # END for each partition to be added 67 return self 68 # END add mode 69 70 raise AssertionError( "Invalid mode given: %i" % mode )
71 72 #} END partition handling 73 74 #{ Member Editing 75
76 - def _toMemberObj( self, member ):
77 """Convert member to a valid member object ( MObject, DagPath or Plug )""" 78 memberobj = member 79 80 if isinstance( member, nt.DagNode ): 81 memberobj = member.dagPath() 82 elif isinstance( member, nt.DependNode ): 83 memberobj = member.object() 84 85 return memberobj
86 87
88 - def _forceMembership( self, member, component, is_single_member, ignore_failure ):
89 """Search all sets connected to our partitions 90 for intersecting members and remove them. 91 Finally dd the members in question to us again 92 93 :param member: can be selection list or MObject, MDagPath, MPlug 94 :return: self if everything is fine""" 95 for partition in self.partitions(): 96 for otherset in partition.sets(): 97 if is_single_member: 98 otherset.removeMember( member, component = component ) 99 else: 100 otherset.removeMembers( otherset.intersection( member, sets_are_members = True ) ) 101 # END single member handling 102 # END for each set in partition 103 # END for each partition 104 105 # finally add the member to our set once more - now it should work 106 # do not risk recursion though by setting everything to ignore errors 107 if isinstance( member, api.MSelectionList ): 108 return self._addRemoveMembers( member, self.kAdd, ignore_failure ) 109 else: 110 return self._addRemoveMember( member, component, self.kAdd, ignore_failure )
111 112
113 - def _checkMemberAddResult( self, member, component, mode, ignore_failure, is_single_member ):
114 """Check whether the given member has truly been added to our set 115 and either force membership or raise and exception 116 117 :param is_single_member: if True, member can safely be assumed to be a single member, 118 this speeds up operations as we do not have to use multi-member tests""" 119 if mode in ( self.kAdd, self.kAddForce ): 120 # do we have to check the result ? 121 if mode == self.kAdd and ignore_failure: 122 return self 123 124 # check result - member can be MObject, MDagPath, plug, or selection list 125 numMatches = 1 126 if isinstance( member, api.MSelectionList ): 127 numMatches = member.length() 128 129 # CHECK MEMBERS 130 not_all_members_added = True 131 if is_single_member: 132 not_all_members_added = not self.isMember( member, component = component ) 133 else: 134 not_all_members_added = self.intersection( member, sets_are_members = True ).length() != numMatches 135 136 if not_all_members_added: 137 if mode == self.kAddForce: 138 return self._forceMembership( member, component, is_single_member, ignore_failure ) 139 140 # if we are here, we do not ignore failure, and raise 141 raise ConstraintError( "At least some members of %r could not be added to %r due to violation of exclusivity constraint" % (member,self) ) 142 143 # END if added members are not yet available 144 # END if mode is add or forced add 145 return self
146 147
148 - def _addRemoveMember( self, member, component, mode, ignore_failure ):
149 """Add or remove the member with undo support 150 151 :param mode: kRemove or kAdd""" 152 memberobj = self._toMemberObj( member ) 153 154 op = undo.GenericOperation() 155 mfninst = self._mfncls( self._apiobj ) 156 doitfunc = mfninst.addMember 157 undoitfunc = mfninst.removeMember 158 159 # for dag paths, append empty component mobjects 160 args = [ memberobj ] 161 if isinstance( memberobj, api.MDagPath ): # add component ( default None ) 162 args.append( component ) 163 # END component handling 164 165 # swap functions if we remove the node 166 if mode == ObjectSet.kRemove: 167 tmp = undoitfunc 168 undoitfunc = doitfunc 169 doitfunc = tmp 170 # END variable switching 171 172 op.setDoitCmd( doitfunc, *args ) 173 op.setUndoitCmd( undoitfunc, *args ) 174 op.doIt( ) 175 176 return self._checkMemberAddResult( member, component, mode, ignore_failure, True )
177
178 - def _addRemoveMembers( self, members, mode, ignore_failure ):
179 """Add or remove the members to the set 180 181 :param mode: kRemove or kAdd or kAddForce""" 182 sellist = nt.toSelectionList( members ) # handles 'member is SelectionList' case ! 183 184 lsellist = sellist.length() 185 if not lsellist: 186 return self 187 188 # if there is only one member, use our single member function 189 # as it will be faster when checking for partition constraints 190 if lsellist == 1: 191 return self._addRemoveMember( it.iterSelectionList( sellist, asNode = 0 ).next(), api.MObject(), mode, ignore_failure ) 192 193 # prepare operation 194 mfninst = self._mfncls( self._apiobj ) 195 doitfunc = mfninst.addMembers 196 undoitfunc = mfninst.removeMembers 197 198 # swap functions if we remove the node 199 if mode == ObjectSet.kRemove: 200 tmp = undoitfunc 201 undoitfunc = doitfunc 202 doitfunc = tmp 203 204 # IMPORTANT: If one member of sellist is not in the set, the operation 205 # will *silently* ( WTF ??) fail. Hence we have to make sure that 206 # we only even remotely think about trying to remove items which are 207 # actually in the set ! 208 sellist = self.intersection(sellist) 209 # END function swapping 210 211 op = undo.GenericOperation() 212 op.setDoitCmd( doitfunc, sellist ) 213 op.setUndoitCmd( undoitfunc, sellist ) 214 op.doIt() 215 216 return self._checkMemberAddResult( sellist, None, mode, ignore_failure, False )
217 218 @undoable
219 - def clear( self ):
220 """Clear the set so that it will be empty afterwards 221 222 :return: self""" 223 self.removeMembers( self.getMembers() ) 224 return self
225 226 @undoable
227 - def addMember( self, member, component = api.MObject(), force = False, ignore_failure = False ):
228 """Add the item to the set 229 230 :param member: Node, MObject, MDagPath or plug 231 :param force: if True, member ship will be forced by removing the member in question 232 from the other set connected to our partitions 233 :param ignore_failure: if True, a failed add due to partion constraints will result in an 234 exception, otherwise it will be silently ignored. Ignored if if force is True 235 :param component: if member is a dagnode, you can specify a component instance 236 of type component instance ( Single|Double|TripleIndexComponent ) 237 :todo: handle components - currently its only possible when using selection lists 238 :return: self """ 239 mode = self.kAdd 240 if force: 241 mode = self.kAddForce 242 return self._addRemoveMember( member, component, mode, ignore_failure )
243 244 @undoable
245 - def add( self, member_or_members, *args, **kwargs ):
246 """Combined method which takes single or multiple members which are to be added 247 248 :param member_or_members: one of the input types supported by `addMember` and 249 `addMembers` 250 :param kwargs: see `addMember` 251 :note: this method is for convenience only and should not be used to 252 add massive amounts of items""" 253 addfun = None 254 if isinstance(member_or_members, (tuple, list, api.MSelectionList)): 255 addfun = self.addMembers 256 else: 257 addfun = self.addMember 258 # END handle input 259 return addfun(member_or_members, *args, **kwargs)
260 261 @undoable
262 - def removeMember( self, member, component = api.MObject() ):
263 """Remove the member from the set 264 265 :param member: member of the list, for types see `addMember`""" 266 return self._addRemoveMember( member, component, ObjectSet.kRemove, True )
267 268 @undoable
269 - def discard( self, member_or_members, *args, **kwargs ):
270 """Removes a single member or multiple members from the set 271 272 :param member_or_members: any of the types supported by `removeMember` 273 or `removeMembers` 274 :param kwargs: see `removeMember`""" 275 rmfun = None 276 if isinstance(member_or_members, (tuple, list, api.MSelectionList)): 277 rmfun = self.removeMembers 278 else: 279 rmfun = self.removeMember 280 # END handle input 281 return rmfun(member_or_members, *args, **kwargs)
282 283 @undoable
284 - def addMembers( self, nodes, force = False, ignore_failure = False ):
285 """Add items from iterable or selection list as members to this set 286 287 :param nodes: MSelectionList or list of Nodes and Plugs 288 :param force: see `addMember` 289 :param ignore_failure: see `addMember` 290 :return: self """ 291 mode = self.kAdd 292 if force: 293 mode = self.kAddForce 294 return self._addRemoveMembers( nodes, mode, ignore_failure )
295 296 @undoable
297 - def removeMembers( self, nodes ):
298 """Remove items from iterable or selection list from this set 299 300 :param nodes: see `addMembers` 301 :return: self """ 302 return self._addRemoveMembers( nodes, ObjectSet.kRemove, True )
303 304 @undoable
305 - def setMembers( self, nodes, mode, **kwargs ):
306 """Adjust set membership for nodes 307 308 :param nodes: items to handle, supports everything that `addMembers` does 309 :param kwargs: arguments passed to `addMembers` or `removeMembers`""" 310 if mode == self.kReplace: 311 self.clear() 312 mode = self.kAdd 313 # END replace 314 315 if mode == self.kAdd: 316 return self.addMembers( nodes, **kwargs ) 317 318 # remove 319 return self.removeMembers( nodes, **kwargs )
320 321 #} END member editing 322 323 324 #{ Member Query 325
326 - def getMembers( self, flatten = False ):
327 """:return: MSelectionList with members of this set 328 :param flatten: if True, members that are objectSets themselves will be resolved to their 329 respective members 330 :note: the members are ordinary api objects that still need to be wrapped 331 :note: use iterMembers to iterate the members as wrapped Nodes""" 332 sellist = api.MSelectionList() 333 self._mfncls( self._apiobj ).getMembers( sellist, flatten ) 334 return sellist
335
336 - def iterMembers( self, *args, **kwargs ):
337 """Iterate members of this set 338 339 :note: All keywords of iterMembers are supported 340 :note: if 'handlePlugs' is False, the iteration using a filter type will be faster 341 :note: handleComponents will allow component iteration - see the iterator documentation""" 342 return it.iterSelectionList( self.getMembers( ), *args, **kwargs )
343
344 - def isMember( self, obj, component = api.MObject() ):
345 """:return: True if obj is a member of this set 346 :param component: is given, the component must be fully part of the set 347 for the object ( dagNode ) to be considered part of the set 348 :note: all keywords of `it.iterSelectionList` are supported 349 :note: ismember does not appear to be working properly with component assignments. 350 It returns true for components that are not actually in the givne shading group""" 351 if not component.isNull(): 352 return self._mfncls( self._apiobj ).isMember( self._toMemberObj( obj ), component ) 353 return self._mfncls( self._apiobj ).isMember( self._toMemberObj( obj ) )
354 355 #} END member query 356 357 # Aliases 358 members = getMembers 359 360 361 #{ Set Operations 362
363 - class _TmpSet( object ):
364 """Temporary set that will delete itself once its python destructor is called""" 365 __slots__ = "setobj"
366 - def __init__( self, sellist ):
367 dgmod = api.MDGModifier( ) 368 self.setobj = dgmod.createNode( "objectSet" ) 369 dgmod.doIt( ) 370 # add members 371 mfnset = api.MFnSet( self.setobj ) 372 mfnset.addMembers( sellist )
373
374 - def __del__( self ):
375 """Delete our own set upon deletion""" 376 # assure we release members before - otherwise they might be deleted 377 # as well if it is empty sets ! 378 mfnset = api.MFnSet( self.setobj ) 379 mfnset.clear() 380 del( mfnset ) 381 dgmod = api.MDGModifier() 382 dgmod.deleteNode( self.setobj ) 383 dgmod.doIt()
384 385 386 @classmethod
387 - def _toValidSetOpInput( cls, objects, sets_are_members = False ):
388 """Method creating valid input for the union/intersection or difference methods 389 390 :note: it may return a temporary set that will delete itself once the wrapper object 391 is being destroyed 392 :param sets_are_members: see `union` 393 :note: set """ 394 if isinstance( objects, (tuple, list) ): 395 # MOBJECTARRAY OF SETS 396 if not objects: # emty list, return empty mobject array 397 return api.MObjectArray( ) 398 399 if not sets_are_members and isinstance( objects[ 0 ], ObjectSet ): 400 objarray = api.MObjectArray( ) 401 for setNode in objects: 402 objarray.append( setNode._apiobj ) 403 return objarray 404 else: 405 # create selection list from nodes and use a tmpSet 406 sellist = nt.toSelectionList( objects ) 407 return cls._TmpSet( sellist ) 408 # END list handling 409 410 # still here, handle a single object 411 singleobj = objects 412 if isinstance( singleobj, api.MSelectionList ): # Selection List ? 413 return cls._TmpSet( singleobj ) 414 415 if not sets_are_members and isinstance( singleobj, ObjectSet ): # Single Object Set ? 416 return singleobj.object() 417 418 if isinstance( singleobj, cls._TmpSet ): # single set object 419 return singleobj.setobj 420 421 if isinstance( singleobj, api.MObject ) and singleobj.hasFn( api.MFn.kSet ): # MObject object set ? 422 return singleobj 423 424 # assume best for MObject arrays - usually we pass it in ourselves 425 if isinstance( singleobj, api.MObjectArray ): 426 return singleobj 427 428 # Can be Node, MDagPath or plug or MObject ( not set ) 429 return cls._toValidSetOpInput( ( singleobj, ), sets_are_members = sets_are_members ) # will create a tmpset then 430 431 raise TypeError( "Type InputObjects for set operation ( %r ) was not recognized" % objects )
432 433
434 - def _applySetOp( self, objects, opid, **kwargs ):
435 """Apply the set operation with the given id""" 436 # have to do it in steps to assure our temporary set will be deleted after 437 # the operation has finished 438 obj = fobj = self._toValidSetOpInput( objects, **kwargs ) 439 outlist = api.MSelectionList() 440 if isinstance( obj, self._TmpSet ): 441 fobj = obj.setobj # need to keep reference to _TmpSet until it was used 442 443 mfnset = self._mfncls( self._apiobj ) 444 if opid == "union": 445 mfnset.getUnion( fobj, outlist ) 446 elif opid == "intersection": 447 mfnset.getIntersection( fobj, outlist ) 448 else: 449 raise AssertionError( "Invalid Set Operation: %s" % opid ) 450 451 return outlist
452 453 @classmethod
454 - def tmpSet( cls, objects, sets_are_members = False ):
455 """ 456 :return: temporary set that will delete itself once it's reference count 457 reaches 0. Use rval.setobj to access the actual set, as the returned object is 458 just a hanlde to it. The handle is a valid input to the set functions as well 459 :param objects: see `union` 460 :param sets_are_members: see `union` 461 :note: useful if you want to use the set member union, intersection or substraction 462 methods efficiently on many sets in a row - these internally operate on a set, thus 463 it is faster to use them with another set from the beginning to prevent creation of intermediate 464 sets""" 465 return cls._toValidSetOpInput( objects, sets_are_members = sets_are_members )
466
467 - def getUnion( self, objects, sets_are_members = False ):
468 """Create a union of the given items with the members of this set 469 470 :param objects: an ObjectSet, an MObject of an object set, a list of ObjectSets 471 or a list of wrapped Objects or an MSelectionList or a single wrapped object . 472 If you have objects in a list as well as sets 473 themselves, objects must come first as the operation will fail otherwise. 474 :param sets_are_members: if True, objects can contain sets, but they should not be treated 475 as sets to apply the set operation with, they should simply be members of this set, and 476 thus need to be wrapped into a tmp set as well 477 :return: MSelectionList of all objects of self and objects """ 478 return self._applySetOp( objects, "union", sets_are_members = sets_are_members )
479
480 - def getIntersection( self, objects, sets_are_members = False ):
481 """As `union`, but returns the intersection ( items in common ) of this 482 set with objects 483 484 :param objects: see `union` 485 :param sets_are_members: see `union` 486 :return: MSelectionList of objects being in self and in objects""" 487 return self._applySetOp( objects, "intersection", sets_are_members = sets_are_members )
488
489 - def getDifference( self, objects, sets_are_members = False ):
490 """return the result of ``self minus objects``, thus objects will be substracted from our obejcts 491 492 :param objects: see `union` 493 :param sets_are_members: see `union` 494 :return: MSelectionList containing objects of self not being in objects list""" 495 # have to do the intersections individually and keep them 496 intersections = list() 497 obj = fobj = self._toValidSetOpInput( objects, sets_are_members = sets_are_members ) 498 outlist = api.MSelectionList() 499 if isinstance( obj, self._TmpSet ): 500 fobj = obj.setobj # need to keep reference to _TmpSet until it was used 501 502 # either we have a single _tmpSet, or a list of sets 503 #if not isinstance( fobj , ( tuple, list ) ): 504 if not hasattr( fobj, '__iter__' ): 505 fobj = [ fobj ] 506 507 for item in fobj: 508 intersections.append( self.intersection( item ) ) 509 510 # remove intersecting members temporarily 511 for its in intersections: 512 self.removeMembers( its ) 513 514 difference = self.getMembers() 515 516 # add members again 517 for its in intersections: 518 self.addMembers( its ) 519 520 return difference
521
522 - def iterUnion( self, setOrSetsOrObjects, **kwargs ):
523 """As union, but returns an iterator 524 525 :param kwargs: passed to it.iterSelectionList""" 526 return it.iterSelectionList( self.union( setOrSetsOrObjects ), **kwargs )
527
528 - def iterIntersection( self, setOrSetsOrObjects, **kwargs ):
529 """As intersection, but returns an iterator 530 531 :param kwargs: passed to it.iterSelectionList""" 532 return it.iterSelectionList( self.intersection( setOrSetsOrObjects ), **kwargs )
533
534 - def iterDifference( self, setOrSetsOrObjects, **kwargs ):
535 """As difference, but returns an iterator 536 537 :param kwargs: passed to it.iterSelectionList""" 538 return it.iterSelectionList( self.difference( setOrSetsOrObjects ), **kwargs )
539 540 #} END set operations 541 542 # aliases 543 union = getUnion 544 intersection = getIntersection 545 difference = getDifference 546 547 #{ Operators 548 __or__ = union 549 __add__ = union 550 __sub__ = difference 551 __and__ = intersection 552 #} END operators 553 554 #{ Protocols
555 - def __len__(self):
556 """ 557 :warn: This method is possibly slow as it will retrieve all members 558 just to get the size of the set. Don't use it directly if you like performance""" 559 return len(self.getMembers())
560
561 - def __iter__(self):
562 return self.getMembers().mtoIter()
563
564 - def __contains__( self, obj ):
565 """:return: True if the given obj is member of this set""" 566 return self.isMember( obj )
567 #} END protocols
568 569 570 571 -class ShadingEngine:
572 """Provides specialized methods able to deal better with shaders 573 than the default implementation. 574 575 :todo: Force exclusivity must be a little more elaborate - this could be overwritten 576 and reimplemented to take care of the details""" 577 578 __metaclass__ = typ.MetaClassCreatorNodes
579
580 581 582 -class Partition:
583 """Deal with common set <-> partition interactions""" 584 __metaclass__ = typ.MetaClassCreatorNodes 585 586 #{ Set Membership 587 588 @undoable
589 - def _addRemoveMember( self, objectset, mode ):
590 sets = objectset 591 if isinstance( objectset, ObjectSet ): 592 sets = [ objectset ] 593 594 for oset in sets: 595 oset.setPartition( self, mode ) 596 # END for each set to add/remove 597 598 return self
599 600
601 - def addMember( self, objectset ):
602 """Add the given objectset or list of sets to the partition 603 604 :param objectset: one or multiple object sets 605 :return: self allowing chained calls""" 606 return self._addRemoveMember( objectset, ObjectSet.kAdd )
607 608 add = addMember 609
610 - def removeMember( self, objectset ):
611 """Remove the given objectset from the partition 612 613 :param objectset: one or multiple object sets 614 :return: self allowing chained calls""" 615 return self._addRemoveMember( objectset, ObjectSet.kRemove )
616 617 discard = removeMember 618
619 - def replaceMember( self, objectset ):
620 """Replace existing objectsets with the given one(s) 621 622 :param objectset: one or multiple object sets 623 :return: self allowing chained calls)""" 624 return self._addRemoveMember( objectset, ObjectSet.kReplace )
625 626 @undoable
627 - def clear( self ):
628 """remove all members from this partition 629 630 :return: self""" 631 for m in self.getMembers(): 632 self.removeMember( m ) 633 634 return self
635
636 - def getMembers( self ):
637 """:return: sets being member of this partition 638 :note: have to filter the members as there might be non-set connections 639 in referenced environments""" 640 out = list() 641 for plug in self.st.minputs(): 642 node = plug.mwrappedNode() 643 if not node.hasFn( api.MFn.kSet ): 644 continue 645 out.append( node ) 646 # END for each plug in set connections 647 return out
648 649 650 #}END set membership 651 652 #{ Name Remapping 653 addSets = addMember 654 removeSets = removeMember 655 replaceSets = replaceMember 656 sets = getMembers 657 members = getMembers 658 #} END name remapping 659 660 #{ Protocols
661 - def __len__(self):
662 return len(self.st.minputs())
663
664 - def __iter__(self):
665 for s in self.getMembers(): 666 yield s
667 # END for each member (ObjectSet) 668 669 #} END protocols 670