Python API 2.0 Reference
scripted/locatorHelperShape.py
1 ###############################################################################
2 #
3 # locatorHelperShape.py
4 #
5 # A plug-in locator with configurable shape representations. The locator
6 # supports reading in a library of shape representations serialized as
7 # JSON files, that populates its 'shapes' enumeration attribute.
8 #
9 # The plug-in provides a command to serialize combinations of meshes
10 # and curves into the library folder.
11 #
12 ## Usage ##
13 #
14 # 1. Create a locatorHelperShape node. By default, the node will have
15 # a default representation consisting of a mesh cube and
16 # diamond shaped curve.
17 #
18 # import maya.cmds as cmds
19 # cmds.createNode('locatorHelperShape')
20 #
21 # 2. Populate the locatorHelperShape library.
22 #
23 # By default, the library is located at $(PLUGIN_PATH)/library/.
24 # Alternately, you can set the library location using the
25 # 'MAYA_LOCATOR_HELPER_SHAPE_LIB' environment variable. This must be set
26 # prior to loading the plug-in.
27 #
28 # import maya.cmds as cmds
29 #
30 # # Create a sphere representation
31 # sph = cmds.polySphere(name='sphere')[0]
32 # cmds.locatorHelperShapeExport(sph, force=True, name=sph)
33 #
34 # # Create a representation with multiple shapes
35 # c = cmds.circle(name='circle')[0]
36 # s = cmds.nurbsSquare(name='square')[0]
37 # cmds.locatorHelperShapeExport(s, c, force=True, name='combo')
38 #
39 # # Can alternately create using the selection list
40 # cmds.select(c, s, r=True)
41 # cmds.locatorHelperShapeExport(force=True, name='combo')
42 #
43 # # Can adjust the number of samples to take from the curve
44 # # when tessellating them into line segments. Default is 50.
45 # cmds.locatorHelperShapeExport(force=True, name='combo', curveSamples=100)
46 #
47 # # Shapes are exported in object space. Freeze
48 # # transforms as necessary to get the desired output.
49 # cmds.scale(1, 2, 1, c)
50 # cmds.makeIdentity(c, apply=True,t=1,r=1,s=1,n=0,pn=1)
51 # cmds.locatorHelperShapeExport(c, force=True, name='oval')
52 #
53 # 3. Reload the plug-in to see the new shapes. You may need to
54 # reload the Attribute Editor to view the updated enumeration.
55 # You can reload the AE by:
56 #
57 # a. Select a locatorHelperShape node.
58 # b. In the AE, navigate the AE menu: 'Show > Set Global View > Lookdev'
59 # b. Restore the AE global view: 'Show > Set Global View > Default'
60 #
61 ## Overview ##
62 #
63 # The plug-in can be broken down into the following components:
64 #
65 # 1. helperShapeRep / helperShapeRepItem
66 #
67 # Defines the shape representations and are responsible for
68 # serialization and generation of the vertex/index buffers
69 # for display.
70 #
71 # 2. helperShapeExportCmd
72 #
73 # Defines the command to convert a list or selection of meshes
74 # and/or curves into helperShapeRepItems for serialization.
75 #
76 # 3. helperShapeNode
77 #
78 # Defines the plug-in locator node.
79 #
80 # 4. helperShapeSubSceneOverride
81 #
82 # Defines the VP2 override implementation for rendering the
83 # helperShapeNode.
84 #
85 import os
86 import sys
87 import ctypes
88 import json
89 
90 import maya.cmds as cmds
91 import maya.mel as mel
92 import maya.api.OpenMaya as om
93 import maya.api.OpenMayaRender as omr
94 import maya.api.OpenMayaUI as omui
95 
96 maya_useNewAPI = True
97 
98 ###############################################################################
99 #
100 # Constants
101 #
102 DEFAULT_COLOR = (0.0, 0.7, 0.15)
103 DEFAULT_LINE_WIDTH = 1.0
104 
105 HELPER_SHAPE_SELECTION_MASK = 'locatorHelperShapeSelection'
106 HELPER_SHAPE_SELECTION_PRIORITY = 9
107 
108 # SubScene Consolidation support
109 ENABLE_CONSOLIDATION = True
110 consolidation_func = getattr(omr.MRenderItem, 'wantSubSceneConsolidation', None)
111 SUPPORTS_CONSOLIDATION = callable(consolidation_func)
112 WANT_CONSOLIDATION = SUPPORTS_CONSOLIDATION and ENABLE_CONSOLIDATION
113 
114 ###############################################################################
115 #
116 # Utility functions
117 #
118 def isclose(a, b, rel_tol=1e-9, abs_tol=0.0):
119  return abs(a-b) <= max( rel_tol * max(abs(a), abs(b)), abs_tol )
120 
121 def isclose_tuple(a, b, rel_tol=1e-9, abs_tol=0.0):
122  for (aa, bb) in zip(a, b):
123  if not isclose(aa, bb):
124  return False
125  return True
126 
127 
128 ###############################################################################
129 #
130 # helperShapeRep / helperShapeRepItem
131 #
132 # The helperShapeRep* classes define the abstraction for the shapes that
133 # helperShapeNode can display. These classes are responsible for:
134 #
135 # 1. Owning the data that defines the shapes.
136 # 2. Computing the associated vertex and index buffers.
137 # 3. Serializing/Deserializing to Python dictionaries.
138 #
139 class helperShapeRepItemType():
140  ''' helperShape representation item type definition. '''
141  kMesh = None
142  kCurve = None
143 
144  def __init__(self, index, name):
145  self.index = index
146  self.name = name
147 
148  def __eq__(self, other):
149  return self.index == other.index
150 
151  def __ne__(self, other):
152  return not(self == other)
153 
154  def __int__(self):
155  return self.index
156 
157  def __str__(self):
158  return self.name
159 
160 # Populate our helperShapeRepItem type variables.
161 helperShapeRepItemType.kMesh = helperShapeRepItemType(0, 'mesh')
162 helperShapeRepItemType.kCurve = helperShapeRepItemType(1, 'curve')
163 
164 class helperShapeRepItem(object):
165  ''' Abstract base class for a helperShape representation item. '''
166  kTypeKey = 'type'
167 
168  def __init__(self, type):
169  ''' Constructor. '''
170  self.type = type
171 
172  @classmethod
173  def deserialize_type(cls, data):
174  ''' Deserialize the item type from a Python dictionary '''
175  if data[cls.kTypeKey] == helperShapeRepItemType.kMesh.name:
176  return helperShapeRepItemType.kMesh
177  elif data[cls.kTypeKey] == helperShapeRepItemType.kCurve.name:
178  return helperShapeRepItemType.kCurve
179  raise RuntimeError('Invalid helperShapeRepItemType: {}'.format(data[cls.kTypeKey]))
180 
181  def deserialize(self, data):
182  ''' Deserialize the item from a Python dictionary '''
183  raise NotImplementedError
184 
185  def serialize(self, data):
186  ''' Serialize the item to a Python dictionary '''
187  if self.type == helperShapeRepItemType.kMesh:
188  data[self.kTypeKey] = helperShapeRepItemType.kMesh.name
189  elif self.type == helperShapeRepItemType.kCurve:
190  data[self.kTypeKey] = helperShapeRepItemType.kCurve.name
191 
192  def vertex_buffer(self):
193  ''' Returns an MVertexBufferArray containing the vertex buffers for this item. '''
194  raise NotImplementedError
195 
196  def wire_index_buffer(self):
197  ''' Returns an MIndexBuffer containing the index buffer for wireframe draw. '''
198  raise NotImplementedError
199 
200  def shaded_index_buffer(self):
201  ''' Returns an MIndexBuffer containing the index buffer for shaded draw. '''
202  raise NotImplementedError
203 
204  def clear_buffers(self):
205  ''' Clears cached vertex/index buffers for this item. '''
206  raise NotImplementedError
207 
208  def bounding_box(self):
209  ''' Returns an MBoundingBox for this item. '''
210  raise NotImplementedError
211 
212  @classmethod
213  def _create_index_buffer(cls, indices):
214  ''' Returns an MIndexBuffer populated with the given indices. '''
215  num_indices = len(indices)
216  buffer = omr.MIndexBuffer(omr.MGeometry.kUnsignedInt32)
217  address = buffer.acquire(num_indices, True)
218  data = ((ctypes.c_uint)*num_indices).from_address(address)
219  for (i,val) in enumerate(indices):
220  data[i] = val
221  buffer.commit(address)
222  return buffer
223 
224  @classmethod
225  def _create_vertex_buffer(cls, desc, vertex_data):
226  ''' Returns an MVertexBuffer populated with the given vertex data. '''
227  num_vertices = len(vertex_data)
228  buffer = omr.MVertexBuffer(desc)
229  address = buffer.acquire(num_vertices, True)
230  if desc.dataType == omr.MGeometry.kFloat and desc.dimension == 3:
231  data = ((ctypes.c_float * 3)*num_vertices).from_address(address)
232  for (i,val) in enumerate(vertex_data):
233  data[i][0] = val[0]
234  data[i][1] = val[1]
235  data[i][2] = val[2]
236  buffer.commit(address)
237  else:
238  # This plug-in only requires working with float3 data.
239  raise NotImplementedError
240  return buffer
241 
242 class helperShapeRepItemMesh(helperShapeRepItem):
243  ''' Defines a mesh based helperShape representation item. '''
244 
245  kVerticesKey = 'vertices'
246  kNormalsKey = 'normals'
247  kWireIndicesKey = 'wireIndices'
248  kShadedIndicesKey = 'shadedIndices'
249 
250  def __init__(self):
251  ''' Constructor. '''
252  super(helperShapeRepItemMesh, self).__init__(helperShapeRepItemType.kMesh)
253 
254  self.position_desc = omr.MVertexBufferDescriptor('', omr.MGeometry.kPosition, omr.MGeometry.kFloat, 3)
255  self.normal_desc = omr.MVertexBufferDescriptor('', omr.MGeometry.kNormal, omr.MGeometry.kFloat, 3)
256 
257  self.vertices = None
258  self.normals = None
259 
260  self.wire_indices = None
261  self.shaded_indices = None
262 
263  self.position_buffer = None
264  self.normal_buffer = None
265  self.index_buffer_wire = None
266  self.index_buffer_shaded = None
267 
268  def deserialize(self, data):
269  ''' Deserialize the item from a Python dictionary '''
270  self.vertices = data[self.kVerticesKey]
271  self.normals = data[self.kNormalsKey]
272  self.wire_indices = data[self.kWireIndicesKey]
273  self.shaded_indices = data[self.kShadedIndicesKey]
274 
275  def serialize(self, data):
276  ''' Serialize the item to a Python dictionary '''
277  super(helperShapeRepItemMesh, self).serialize(data)
278  data[self.kVerticesKey] = self.vertices
279  data[self.kNormalsKey] = self.normals
280  data[self.kWireIndicesKey] = self.wire_indices
281  data[self.kShadedIndicesKey] = self.shaded_indices
282 
283  def vertex_buffer(self):
284  ''' Returns an MVertexBufferArray containing the vertex buffers for this item. '''
285  if not self.position_buffer:
286  self.position_buffer = self._create_vertex_buffer(self.position_desc, self.vertices)
287  if not self.normal_buffer:
288  self.normal_buffer = self._create_vertex_buffer(self.normal_desc, self.normals)
289  vertex_buffer = omr.MVertexBufferArray()
290  vertex_buffer.append(self.position_buffer, 'positions')
291  vertex_buffer.append(self.normal_buffer, 'normals')
292  return vertex_buffer
293 
294  def wire_index_buffer(self):
295  ''' Returns an MIndexBuffer containing the index buffer for wireframe draw. '''
296  if not self.index_buffer_wire:
297  self.index_buffer_wire = self._create_index_buffer(self.wire_indices)
298  return self.index_buffer_wire
299 
300  def shaded_index_buffer(self):
301  ''' Returns an MIndexBuffer containing the index buffer for shaded draw. '''
302  if not self.index_buffer_shaded:
303  self.index_buffer_shaded = self._create_index_buffer(self.shaded_indices)
304  return self.index_buffer_shaded
305 
306  def clear_buffers(self):
307  ''' Clears cached vertex/index buffers for this item. '''
308  self.position_buffer = None
309  self.normal_buffer = None
310  self.index_buffer_wire = None
311  self.index_buffer_shaded = None
312 
313  def bounding_box(self):
314  ''' Returns an MBoundingBox for this item. '''
315  box = om.MBoundingBox()
316  for v in self.vertices:
317  box.expand(om.MPoint(v[0], v[1], v[2]))
318  return box
319 
320 class helperShapeRepItemCurve(helperShapeRepItem):
321  ''' Defines a curve based helperShape representation item. '''
322 
323  kVerticesKey = 'vertices'
324  kIndicesKey = 'indices'
325 
326  def __init__(self):
327  ''' Constructor. '''
328  super(helperShapeRepItemCurve, self).__init__(helperShapeRepItemType.kCurve)
329 
330  self.position_desc = omr.MVertexBufferDescriptor('', omr.MGeometry.kPosition, omr.MGeometry.kFloat, 3)
331 
332  self.vertices = None
333  self.indices = None
334 
335  self.position_buffer = None
336  self.index_buffer = None
337 
338  def deserialize(self, data):
339  ''' Deserialize the item from a Python dictionary '''
340  self.vertices = data[self.kVerticesKey]
341  self.indices = data[self.kIndicesKey]
342 
343  def serialize(self, data):
344  ''' Serialize the item to a Python dictionary '''
345  super(helperShapeRepItemCurve, self).serialize(data)
346  data[self.kVerticesKey] = self.vertices
347  data[self.kIndicesKey] = self.indices
348 
349  def vertex_buffer(self):
350  ''' Returns an MVertexBufferArray containing the vertex buffers for this item. '''
351  if not self.position_buffer:
352  self.position_buffer = self._create_vertex_buffer(self.position_desc, self.vertices)
353  vertex_buffer = omr.MVertexBufferArray()
354  vertex_buffer.append(self.position_buffer, 'positions')
355  return vertex_buffer
356 
357  def wire_index_buffer(self):
358  ''' Returns an MIndexBuffer containing the index buffer for wireframe draw. '''
359  if not self.index_buffer:
360  self.index_buffer = self._create_index_buffer(self.indices)
361  return self.index_buffer
362 
363  def clear_buffers(self):
364  ''' Clears cached vertex/index buffers for this item. '''
365  self.position_buffer = None
366  self.index_buffer = None
367 
368  def bounding_box(self):
369  ''' Returns an MBoundingBox for this item. '''
370  box = om.MBoundingBox()
371  for v in self.vertices:
372  box.expand(om.MPoint(v[0], v[1], v[2]))
373  return box
374 
375 class helperShapeRepItemCube(helperShapeRepItemMesh):
376  ''' A sample default cube mesh representation item. '''
377  def __init__(self):
378  super(helperShapeRepItemCube, self).__init__()
379 
380  # Exploded vertex layout (unshared normals)
381  self.vertices = [
382  #top
383  (-0.5, 0.5,-0.5),
384  (-0.5, 0.5, 0.5),
385  ( 0.5, 0.5, 0.5),
386  ( 0.5, 0.5,-0.5),
387 
388  #left
389  (-0.5,-0.5,-0.5),
390  (-0.5,-0.5, 0.5),
391  (-0.5, 0.5, 0.5),
392  (-0.5, 0.5,-0.5),
393 
394  #bottom
395  ( 0.5,-0.5, 0.5),
396  (-0.5,-0.5, 0.5),
397  (-0.5,-0.5,-0.5),
398  ( 0.5,-0.5,-0.5),
399 
400  #right
401  ( 0.5,-0.5,-0.5),
402  ( 0.5, 0.5,-0.5),
403  ( 0.5, 0.5, 0.5),
404  ( 0.5,-0.5, 0.5),
405 
406  #front
407  (-0.5,-0.5, 0.5),
408  ( 0.5,-0.5, 0.5),
409  ( 0.5, 0.5, 0.5),
410  (-0.5, 0.5, 0.5),
411 
412  #back
413  (-0.5,-0.5,-0.5),
414  (-0.5, 0.5,-0.5),
415  ( 0.5, 0.5,-0.5),
416  ( 0.5,-0.5,-0.5)
417  ]
418 
419  self.normals = [
420  # top
421  ( 0.0, 1.0, 0.0),
422  ( 0.0, 1.0, 0.0),
423  ( 0.0, 1.0, 0.0),
424  ( 0.0, 1.0, 0.0),
425 
426  #left
427  (-1.0, 0.0, 0.0),
428  (-1.0, 0.0, 0.0),
429  (-1.0, 0.0, 0.0),
430  (-1.0, 0.0, 0.0),
431 
432  #bottom
433  ( 0.0,-1.0, 0.0),
434  ( 0.0,-1.0, 0.0),
435  ( 0.0,-1.0, 0.0),
436  ( 0.0,-1.0, 0.0),
437 
438  #right
439  ( 1.0, 0.0, 0.0),
440  ( 1.0, 0.0, 0.0),
441  ( 1.0, 0.0, 0.0),
442  ( 1.0, 0.0, 0.0),
443 
444  #front
445  ( 0.0, 0.0, 1.0),
446  ( 0.0, 0.0, 1.0),
447  ( 0.0, 0.0, 1.0),
448  ( 0.0, 0.0, 1.0),
449 
450  #back
451  ( 0.0, 0.0,-1.0),
452  ( 0.0, 0.0,-1.0),
453  ( 0.0, 0.0,-1.0),
454  ( 0.0, 0.0,-1.0)
455  ]
456 
457  num_prim = len(self.vertices)
458  num_index = num_prim * 2
459  self.wire_indices = [None] * num_index
460  i = 0
461  while i < num_index:
462  start_index = i / 8 * 4 # 8 indices per face. 4 verts per face in the vertexBuffer
463  pairs = [(0,1), (1,2), (2,3), (3,0)]
464  for (index, p) in enumerate(pairs):
465  self.wire_indices[i + index*2] = start_index + p[0]
466  self.wire_indices[i + index*2 + 1] = start_index + p[1]
467  i += 8
468 
469  num_prim = len(self.vertices)
470  num_index = num_prim / 4 * 6 # 4 verts per quad, 6 indices per quad (2 tris)
471  self.shaded_indices = [None] * num_index
472  i = 0
473  while i < num_index:
474  start_index = i / 6 * 4
475  tris = [(0,1,2), (0,2,3)]
476  for (index, t) in enumerate(tris):
477  self.shaded_indices[i + index*3] = start_index + t[0]
478  self.shaded_indices[i + index*3 + 1] = start_index + t[1]
479  self.shaded_indices[i + index*3 + 2] = start_index + t[2]
480  i += 6
481 
482 class helperShapeRepItemDiamond(helperShapeRepItemCurve):
483  ''' A sample default diamond curve representation item. '''
484  def __init__(self):
485  super(helperShapeRepItemDiamond, self).__init__()
486  self.vertices = [
487  ( 1.0, 0.0, 0.0),
488  ( 0.0, 0.0, 1.0),
489  (-1.0, 0.0, 0.0),
490  ( 0.0, 0.0,-1.0)
491  ]
492  self.indices = [0,1,1,2,2,3,3,0]
493 
494 class helperShapeRep(object):
495  '''
496  Defines the helperShape representation class.
497  A representation consists of a list of representation items.
498  '''
499 
500  kRenderItemsKey = 'renderItems'
501 
502  def __init__(self):
503  self.name = None
504  self.items = []
505 
506  def deserialize(self, file_path):
507  ''' Deserialize the representation from JSON. '''
508  self.name = os.path.splitext(os.path.basename(file_path))[0]
509  with open(file_path, 'r') as f:
510  data = json.load(f)
511 
512  # Validate data
513  if not isinstance(data, dict):
514  raise RuntimeError('Invalid JSON helperShape file: %s'.format(file_path))
515  if not self.kRenderItemsKey in data or not isinstance(data[self.kRenderItemsKey], list):
516  raise RuntimeError('Invalid [renderItems] key in file: %s'.format(file_path))
517 
518  # Read data
519  for item_data in data[self.kRenderItemsKey]:
520  type = helperShapeRepItem.deserialize_type(item_data)
521  if type == helperShapeRepItemType.kMesh:
522  item = helperShapeRepItemMesh()
523  elif type == helperShapeRepItemType.kCurve:
524  item = helperShapeRepItemCurve()
525  else:
526  continue
527  item.deserialize(item_data)
528  self.items.append(item)
529 
530  def serialize(self, file_path):
531  ''' Serialize the representation to JSON. '''
532  data = {}
533  data[self.kRenderItemsKey] = []
534  for item in self.items:
535  item_data = {}
536  item.serialize(item_data)
537  data[self.kRenderItemsKey].append(item_data)
538  with open(file_path, 'w') as f:
539  json.dump( data, f, sort_keys=True, indent=4, separators=[',', ': '] )
540 
541 class helperShapeRepDefault(helperShapeRep):
542  '''
543  A sample default helper representation consisting of a cube mesh
544  and diamond curve.
545  '''
546  def __init__(self):
547  super(helperShapeRepDefault, self).__init__()
548  self.name = 'default'
549  self.items.append( helperShapeRepItemCube() )
550  self.items.append( helperShapeRepItemDiamond() )
551 
552 
553 ###############################################################################
554 #
555 # helperShapeExportCmd (MPxCommand)
556 #
557 class helperShapeExportCmd(om.MPxCommand):
558  ''' helperShapeExportCmd command class. '''
559  name = 'locatorHelperShapeExport'
560 
561  kNameFlag = 'n'
562  kNameFlagLong = 'name'
563  kCurveSamplesFlag = 'cs'
564  kCurveSamplesFlagLong = 'curveSamples'
565  kForceFlag = 'f'
566  kForceFlagLong = 'force'
567 
568  def __init__(self):
569  ''' Constructor. '''
570  om.MPxCommand.__init__(self)
571 
572  self.name = None
573  self.curve_samples = 50
574  self.force = False
575 
576  @staticmethod
577  def creator():
578  ''' Creator method. '''
579  return helperShapeExportCmd()
580 
581  @staticmethod
582  def new_syntax():
583  ''' Command syntax definition. '''
584  syntax = om.MSyntax()
585  syntax.addFlag(helperShapeExportCmd.kNameFlag, helperShapeExportCmd.kNameFlagLong, om.MSyntax.kString)
586  syntax.addFlag(helperShapeExportCmd.kCurveSamplesFlag, helperShapeExportCmd.kCurveSamplesFlagLong, om.MSyntax.kUnsigned)
587  syntax.addFlag(helperShapeExportCmd.kForceFlag, helperShapeExportCmd.kForceFlagLong)
588  syntax.setObjectType(om.MSyntax.kStringObjects, 0)
589  return syntax
590 
591  @classmethod
592  def valid_name(cls, name):
593  '''
594  Returns true if the name does not already exist in the
595  shape library path. False otherwise.
596  '''
597  for (root, dirs, files) in os.walk(helperShapeNode.shapes_lib_path):
598  for f in files:
599  (file_name, file_ext) = os.path.splitext(f)
600  if not file_ext == '.json':
601  continue
602  if name == file_name:
603  return False
604  return True
605 
606  def get_objects(self, parser):
607  '''
608  Filters the command object list for meshes/curve shapes.
609  If the command object list is empty, filters the selection
610  list for meshes/curve shapes.
611 
612  Returns the meshes/curve shapes in a list of MObjects.
613  '''
614  obj_strs = parser.getObjectStrings()
615  sel_list = om.MSelectionList()
616  if len(obj_strs):
617  for o in obj_strs:
618  sel_list.add(o)
619  else:
620  sel_list = om.MGlobal.getActiveSelectionList()
621 
622  objects = []
623  for i in xrange(sel_list.length()):
624  try:
625  path = sel_list.getDagPath(i)
626  except TypeError:
627  continue
628 
629  multiple_shapes = False
630  try:
631  path.extendToShape()
632  except RuntimeError:
633  multiple_shapes = True
634 
635  if not multiple_shapes:
636  if path.apiType() == om.MFn.kMesh or path.apiType() == om.MFn.kNurbsCurve:
637  objects.append(path.node())
638  else:
639  dag_it = om.MItDag()
640  dag_it.reset(path)
641  while not dag_it.isDone():
642  item = dag_it.getPath()
643  if item.apiType() == om.MFn.kMesh or item.apiType() == om.MFn.kNurbsCurve:
644  objects.append(item.node())
645  dag_it.next()
646  return objects
647 
648  def generate_mesh(self, obj):
649  '''
650  Given a mesh shape MObject, generate a helperShapeRepItemMesh.
651  '''
652  if not obj.apiType() == om.MFn.kMesh:
653  return None
654 
655  # Store raw position/normal mesh data as lists of float3 tuples
656  mesh_fn = om.MFnMesh(obj)
657  raw_vertices_mpts = mesh_fn.getPoints()
658  raw_vertices = [ (p.x, p.y, p.z) for p in raw_vertices_mpts ]
659  raw_normals_mvcs = mesh_fn.getNormals()
660  raw_normals = [ (n.x, n.y, n.z) for n in raw_normals_mvcs ]
661 
662  # Output vertex data arrays
663  vertices = []
664  normals = []
665  wire_indices = []
666  shaded_indices = []
667 
668  # Dictionary with ((position, normal), index) key
669  # value pairs to test for duplicate vertex data.
670  unique_vertices = {}
671 
672  # Mesh vertex IDs to output vertices array
673  mesh_vertex_map = {}
674 
675  # Iterate over each face to generate kTriangles
676  # vertex data.
677  face_it = om.MItMeshPolygon(obj)
678  while not face_it.isDone():
679  count = len(vertices)
680  num_face_verts = face_it.polygonVertexCount()
681 
682  # Create a vertex id to local face vertex id map
683  # to help us correlate face vertex data.
684  mesh_to_face_vertex_map = {}
685  for i in xrange(num_face_verts):
686  mesh_vid = face_it.vertexIndex(i)
687  mesh_to_face_vertex_map[mesh_vid] = i
688 
689  # Iterate over each triangle in this face.
690  (tri_pts, tri_ids) = face_it.getTriangles()
691  num_tris = len(tri_ids) / 3
692  for i in xrange(num_tris):
693  for j in xrange(3):
694  vid = tri_ids[i*3 + j]
695  local_vid = mesh_to_face_vertex_map[vid]
696  nid = face_it.normalIndex(local_vid)
697 
698  # Check if this vertex data combination (position, normal)
699  # already exists in our vertices list. If so, reuse
700  # previous entry. Otherwise, add it and record it in our
701  # unique_vertices lookup map.
702  vertex_data = (raw_vertices[vid], raw_normals[nid])
703  if vertex_data in unique_vertices:
704  vertex_id = unique_vertices[vertex_data]
705  else:
706  vertex_id = count
707  count += 1
708  unique_vertices[vertex_data] = vertex_id
709  vertices.append(vertex_data[0])
710  normals.append(vertex_data[1])
711 
712  # Record the corresponding vertex_id into our
713  # shaded triangle index list.
714  shaded_indices.append(vertex_id)
715 
716  # Record mesh vertex id to vertices list index
717  # for later wireframe generation.
718  if vid in mesh_vertex_map:
719  mesh_vertex_map[vid].append(vertex_id)
720  else:
721  mesh_vertex_map[vid] = [vertex_id,]
722  # BUG: The API docs and function signature suggest this
723  # expects only the self argument, but it errors expecting
724  # 1 argument. Pass in a dummy value of 0 as a workaround.
725  face_it.next(0)
726 
727  # Iterate over each edge to generate wireframe data
728  # in kLines format.
729  edge_it = om.MItMeshEdge(obj)
730  while not edge_it.isDone():
731  mesh_vids = [edge_it.vertexId(0), edge_it.vertexId(1)]
732  vids = [mesh_vertex_map[mesh_vids[0]][0], mesh_vertex_map[mesh_vids[1]][0]]
733  for vid in vids:
734  wire_indices.append(vid)
735  edge_it.next()
736 
737  rep_mesh = helperShapeRepItemMesh()
738  rep_mesh.vertices = vertices
739  rep_mesh.normals = normals
740  rep_mesh.wire_indices = wire_indices
741  rep_mesh.shaded_indices = shaded_indices
742  return rep_mesh
743 
744  def generate_curve(self, obj):
745  '''
746  Given a curve shape MObject, generate a helperShapeRepItemCurve.
747  '''
748  if not obj.apiType() == om.MFn.kNurbsCurve:
749  return None
750  curve_fn = om.MFnNurbsCurve(obj)
751 
752  vertices = []
753  indices = []
754 
755  # Generate vertex positions
756  if curve_fn.degree == 1:
757  # For degree 1 curves, store the CVs.
758  cvs = curve_fn.cvPositions()
759  vertices = [ (v.x, v.y, v.z) for v in cvs ]
760  else:
761  # For degree >= 2 curves, sample the curve.
762  curve_len = curve_fn.length()
763  curve_sample_len = curve_len / (self.curve_samples)
764  curve_len_samples = [ i*curve_sample_len for i in xrange(self.curve_samples) ]
765  for s in curve_len_samples:
766  param = curve_fn.findParamFromLength(s)
767  pt = curve_fn.getPointAtParam(param)
768  vertices.append((pt.x, pt.y, pt.z))
769 
770  # Generate indices
771  for i in xrange(1, len(vertices)):
772  indices.append(i-1)
773  indices.append(i)
774  if curve_fn.form == om.MFnNurbsCurve.kClosed or curve_fn.form == om.MFnNurbsCurve.kPeriodic:
775  indices.append(len(vertices)-1)
776  indices.append(0)
777 
778  rep_curve = helperShapeRepItemCurve()
779  rep_curve.vertices = vertices
780  rep_curve.indices = indices
781  return rep_curve
782 
783  def doIt(self, arg_list):
784  ''' Perform the command. '''
785  parser = om.MArgParser(self.syntax(), arg_list)
786 
787  # Force
788  if parser.isFlagSet(self.kForceFlag):
789  self.force = True
790 
791  # Name
792  if not parser.isFlagSet(self.kNameFlag):
793  raise RuntimeError('The -{}/{} flag must be set.'.format(self.kNameFlag, self.kNameFlagLong))
794  self.name = parser.flagArgumentString(self.kNameFlag, 0)
795  if not self.force and not self.valid_name(self.name):
796  raise RuntimeError('The specified name already exists: {}'.format(self.name))
797 
798  # Curve samples
799  if parser.isFlagSet(self.kCurveSamplesFlag):
800  self.curve_samples = parser.flagArgumentInt(self.kCurveSamplesFlag, 0)
801 
802  # Parse the object list or selection for meshes/curves.
803  objects = self.get_objects(parser)
804 
805  # Generate representations from these objects.
806  rep = helperShapeRep()
807  rep.name = self.name
808  for o in objects:
809  if o.apiType() == om.MFn.kMesh:
810  item = self.generate_mesh(o)
811  elif o.apiType() == om.MFn.kNurbsCurve:
812  item = self.generate_curve(o)
813  else:
814  continue
815  if item:
816  rep.items.append(item)
817 
818  # Serialize the representation.
819  if len(rep.items):
820  lib_path = helperShapeNode.shapes_lib_path
821  if not os.path.exists(lib_path):
822  os.makedirs(lib_path)
823  file_name = self.name + '.json'
824  file_path = os.path.join(lib_path, file_name)
825  rep.serialize(file_path)
826 
827 
828 ###############################################################################
829 #
830 # helperShapeNode (MPxLocatorNode)
831 #
832 class helperShapeNode(omui.MPxLocatorNode):
833  ''' helperShapeNode locator class. '''
834  name = 'locatorHelperShape'
835  id = om.MTypeId( 0x00080041 )
836  drawDbClassification = 'drawdb/subscene/locatorHelperShape'
837  drawRegistrantId = 'locatorHelperShape_SubSceneOverride'
838 
839  aShape = om.MObject()
840  aColor = om.MObject()
841  aDrawOnTop = om.MObject()
842 
843  shapes_lib_path = None
844  shapes = []
845 
846  attrEditorTemplate = '''
847  global proc AElocatorHelperShapeColorNew(string $plug) {
848  AElocatorHelperShapeColorReplace($plug);
849  }
850 
851  global proc AElocatorHelperShapeColorReplace(string $plug) {
852  setUITemplate -pst attributeEditorTemplate;
853  string $parent = `setParent -q`;
854  string $frame = $parent + "|locatorHelperShapeColorFrame";
855  if (`frameLayout -ex $frame`) {
856  deleteUI $frame;
857  }
858  $frame = `frameLayout -l "Colors" -collapse false locatorHelperShapeColorFrame`;
859  string $column = `columnLayout -adj true`;
860  int $colorIds[] = python( "list(set(maya.cmds.getAttr(\\"" + $plug + "\\", mi=True)))" );
861  for ($id in $colorIds) {
862  string $index = "[" + $id + "]";
863  string $childPlug = $plug + $index;
864  string $childName = "color" + $index;
865  attrColorSliderGrp -l $childName -attribute $childPlug;
866  }
867  setUITemplate -ppt;
868  }
869 
870  global proc AElocatorHelperShapeTemplate(string $node) {
871  editorTemplate -beginScrollLayout;
872  editorTemplate -beginLayout "Shape Attributes" -collapse 0;
873  editorTemplate -addControl "shape";
874  editorTemplate -addControl "drawOnTop";
875  editorTemplate -callCustom "AElocatorHelperShapeColorNew"
876  "AElocatorHelperShapeColorReplace"
877  "color";
878  editorTemplate -endLayout;
879  editorTemplate -beginLayout "Locator Attributes";
880  AElocatorCommon $node;
881  editorTemplate -endLayout;
882  AElocatorInclude $node;
883  editorTemplate -addExtraControls;
884  editorTemplate -endScrollLayout;
885  }
886  '''
887 
888  def __init__(self):
889  ''' Constructor. '''
890  omui.MPxLocatorNode.__init__(self)
891 
892  def postConstructor(self):
893  ''' Post-constructor. '''
894  enum_attr_fn = om.MFnEnumAttribute(self.aShape)
895  shape_default = enum_attr_fn.default
896  shape_plug = om.MPlug(self.thisMObject(), self.aShape)
897  shape_plug.setShort(shape_default)
898 
899  @staticmethod
900  def creator():
901  ''' Creator method. '''
902  return helperShapeNode()
903 
904  @classmethod
905  def initialize(cls):
906  ''' Initialize node attribute layout. '''
907  num_attr_fn = om.MFnNumericAttribute()
908  enum_attr_fn = om.MFnEnumAttribute()
909 
910  default_index = 0
911  cls.aShape = enum_attr_fn.create('shape', 'sh')
912  for (i, shape) in enumerate(cls.shapes):
913  enum_attr_fn.addField(shape.name, i)
914  if shape.name == 'default':
915  default_index = i
916  enum_attr_fn.default = default_index
917  enum_attr_fn.internal = True
918 
919  cls.aColor = num_attr_fn.createColor('color', 'cl')
920  num_attr_fn.default = DEFAULT_COLOR
921  num_attr_fn.array = True
922 
923  cls.aDrawOnTop = num_attr_fn.create('drawOnTop', 'dot', om.MFnNumericData.kBoolean, False)
924 
925  cls.addAttribute( cls.aShape )
926  cls.addAttribute( cls.aColor )
927  cls.addAttribute( cls.aDrawOnTop )
928 
929  def compute(self, plug, data):
930  ''' Compute method. '''
931  return None
932 
933  def getShapeSelectionMask(self):
934  ''' helperShape selection mask. '''
935  return om.MSelectionMask(HELPER_SHAPE_SELECTION_MASK)
936 
937  def setInternalValueInContext(self, plug, data_handle, ctx):
938  ''' Callback to set internal attribute values. '''
939  # Use set internal callback to detect when the shape changes
940  # and update the number of color fields accordingly.
941  if plug == self.aShape:
942  # Persist color[0] between shapes.
943  color_plug = om.MPlug(self.thisMObject(), self.aColor)
944  base_color_plug = color_plug.elementByLogicalIndex(0)
945  base_color_obj = base_color_plug.asMObject()
946  base_color = om.MFnNumericData(base_color_obj).getData()
947 
948  shape_id = data_handle.asShort()
949  num_shape_items = len(self.shapes[shape_id].items)
950  data_block = self.forceCache()
951  color_builder = om.MArrayDataBuilder(data_block, self.aColor, num_shape_items)
952  for i in xrange(num_shape_items):
953  child = color_builder.addLast()
954  child.set3Float(base_color[0], base_color[1], base_color[2])
955 
956  color_ovr_hnd = data_block.outputArrayValue(self.aColor)
957  color_ovr_hnd.set(color_builder)
958 
959  # Force an AE refresh.
960  mel.eval('autoUpdateAttrEd;')
961  return False
962 
963  @classmethod
964  def init_shapes(cls, lib_path):
965  '''
966  Populate the list of helperShape representations from
967  the helperShape library path.
968 
969  This must be called prior to node registration since
970  the enumeration attribute to select the active shape
971  depends on the list of shapes.
972  '''
973  # Record lib_path
974  cls.shapes_lib_path = lib_path
975 
976  # Find all *.json files in lib_path
977  shape_files = []
978  if os.path.exists(lib_path):
979  for (root, dirs, files) in os.walk(lib_path):
980  for f in files:
981  (name, ext) = os.path.splitext(f)
982  if ext == '.json':
983  file_path = os.path.join(root, f)
984  shape_files.append(file_path)
985 
986  cls.shapes = []
987  if len(shape_files):
988  for f in shape_files:
989  rep = helperShapeRep()
990  rep.deserialize(f)
991  cls.shapes.append(rep)
992  else:
993  cls.shapes.append(helperShapeRepDefault())
994 
995 
996 ###############################################################################
997 #
998 # helperShapeSubSceneOverride (MPxSubSceneOverride)
999 #
1000 class helperShapeShaderItem(object):
1001  ''' helperShape shader class to help uniquely identify shaders. '''
1002  def __init__(self, shader_type, shader_color, transparent, pre_cb, post_cb):
1003  self.type = shader_type
1004  self.color = shader_color
1005  self.transparent = transparent
1006  self.pre_cb = pre_cb
1007  self.post_cb = post_cb
1008 
1009  def __eq__(self, other):
1010  return (isinstance(other, helperShapeShaderItem) and
1011  self.type == other.type and
1012  isclose_tuple(self.color, other.color) and
1013  self.transparent == other.transparent and
1014  self.pre_cb == other.pre_cb and
1015  self.post_cb == other.post_cb)
1016 
1017  def __ne__(self, other):
1018  return not(self == other)
1019 
1020  def __key(self):
1021  return (self.type, self.color, self.transparent, self.pre_cb, self.post_cb)
1022 
1023  def __hash__(self):
1024  return hash(self.__key())
1025 
1026 class helperShapeShaderCache(object):
1027  ''' helperShape cache of shader instances '''
1028  def __init__(self):
1029  ''' Constructor. Initialize the shader cache. '''
1030  self.cache = {}
1031 
1032  def __del__(self):
1033  ''' Destructor. Clear the shader cache. '''
1034  shader_mgr = omr.MRenderer.getShaderManager()
1035  if not shader_mgr:
1036  return
1037  for (_, shader) in self.cache.items():
1038  shader_mgr.releaseShader(shader)
1039  self.cache = {}
1040 
1041  def get_wire_shader(self, shader_color, transparent, pre_cb=None, post_cb=None):
1042  ''' Return a wire shader with the given parameters from the cache. '''
1043  shader_mgr = omr.MRenderer.getShaderManager()
1044  if not shader_mgr:
1045  return None
1046  shader_type = omr.MShaderManager.k3dThickLineShader
1047  item = helperShapeShaderItem(shader_type, shader_color, transparent, pre_cb, post_cb)
1048  if not item in self.cache:
1049  shader = shader_mgr.getStockShader(shader_type, pre_cb, post_cb)
1050  shader.setParameter('solidColor', shader_color)
1051  shader.setParameter('lineWidth', [DEFAULT_LINE_WIDTH]*2)
1052  shader.setIsTransparent(transparent)
1053  self.cache[item] = shader
1054  return self.cache[item]
1055 
1056  def get_shaded_shader(self, shader_color, transparent, pre_cb=None, post_cb=None):
1057  ''' Return a shaded shader with the given parameters from the cache. '''
1058  shader_mgr = omr.MRenderer.getShaderManager()
1059  if not shader_mgr:
1060  return None
1061  shader_type = omr.MShaderManager.k3dBlinnShader
1062  item = helperShapeShaderItem(shader_type, shader_color, transparent, pre_cb, post_cb)
1063  if not item in self.cache:
1064  shader = shader_mgr.getStockShader(shader_type, pre_cb, post_cb)
1065  shader.setParameter('diffuseColor', shader_color)
1066  shader.setIsTransparent(transparent)
1067  self.cache[item] = shader
1068  return self.cache[item]
1069 
1070 class helperShapeSubSceneOverride(omr.MPxSubSceneOverride):
1071  ''' helperShape viewport subscene override class. '''
1072  kUpdateShaders = 1 << 0
1073  kUpdateGeometry = 1 << 1
1074  kUpdateMatrix = 1 << 2
1075  kUpdateAll = 0xf
1076 
1077  sLeadName = 'lead'
1078  sActiveName = 'active'
1079  sWireName = 'wire'
1080  sShadedName = 'shaded'
1081 
1082  shader_cache = helperShapeShaderCache()
1083 
1084  saved_depth_state = None
1085 
1086  def __init__(self, obj):
1087  ''' Constructor. '''
1088  omr.MPxSubSceneOverride.__init__(self, obj)
1089 
1090  node_fn = om.MFnDependencyNode(obj)
1091  self.node = node_fn.userNode()
1092  self.shape_index = 0
1093  self.shape_colors = []
1094  self.lead_color = (0.0, 0.0, 0.0, 1.0)
1095  self.active_color = (0.0, 0.0, 0.0, 1.0)
1096  self.draw_on_top = False
1097 
1098  self.wire_shaders = []
1099  self.shaded_shaders = []
1100  self.lead_select_shader = None
1101  self.active_select_shader = None
1102 
1103  self.update_mask = self.kUpdateAll
1104 
1105  # Cached instance data (transform_matrix, display_status)
1106  # keyed by instance number.
1107  self.instances = {}
1108 
1109  # Cached lists of render item names sorted by display_status
1110  self.lead_items = []
1111  self.active_items = []
1112  self.wire_items = []
1113  self.shaded_items = []
1114 
1115  @staticmethod
1116  def creator(obj):
1117  ''' Creator method. '''
1118  return helperShapeSubSceneOverride(obj)
1119 
1120  def supportedDrawAPIs(self):
1121  ''' Return the draw APIs supported by this override. '''
1122  return omr.MRenderer.kOpenGL | omr.MRenderer.kDirectX11 | omr.MRenderer.kOpenGLCoreProfile
1123 
1124  def hasUIDrawables(self):
1125  ''' Does this override have UI drawables? '''
1126  return False
1127 
1128  @classmethod
1129  def cleanup(cls):
1130  ''' Cleanup our shader cache. '''
1131  del(cls.shader_cache)
1132  cls.shader_cache = None
1133 
1134  @classmethod
1135  def pre_shader_cb(cls, context, render_items, shader):
1136  ''' Pre shader render callback. Used to facilitate drawOnTop feature. '''
1137  if len(render_items) > 0:
1138  state_mgr = context.getStateManager()
1139  cls.saved_depth_state = state_mgr.getDepthStencilState()
1140  depth_state_desc = omr.MDepthStencilStateDesc()
1141  depth_state_desc.depthEnable = False
1142  disabled_depth_state = state_mgr.acquireDepthStencilState(depth_state_desc)
1143  state_mgr.setDepthStencilState(disabled_depth_state)
1144 
1145  @classmethod
1146  def post_shader_cb(cls, context, render_items, shader):
1147  ''' Post shader render callback. Used to facilitate drawOnTop feature. '''
1148  if len(render_items) > 0:
1149  state_mgr = context.getStateManager()
1150  state_mgr.setDepthStencilState(cls.saved_depth_state)
1151  state_mgr.releaseDepthStencilState(cls.saved_depth_state)
1152  cls.saved_depth_state = None
1153 
1154  def requiresUpdate(self, container, frame_context):
1155  ''' Returns true if the update function should be called. '''
1156  self.requires_update_shaders()
1157  self.requires_update_geometry()
1158  self.requires_update_matrix()
1159  return self.update_mask > 0
1160 
1161  def requires_update_shaders(self):
1162  '''
1163  Checks if any attributes have changed that would affect
1164  this override's shaders. If so, sets the kUpdateShaders
1165  update bit.
1166  '''
1167  shape_plug = om.MPlug(self.node.thisMObject(), helperShapeNode.aShape)
1168  shape_id = shape_plug.asShort()
1169  num_shapes = len(helperShapeNode.shapes[shape_id].items)
1170 
1171  old_colors = self.shape_colors
1172  new_colors = []
1173  color_plug = om.MPlug(self.node.thisMObject(), helperShapeNode.aColor)
1174  for i in xrange(num_shapes):
1175  color_elem = color_plug.elementByLogicalIndex(i)
1176  color_obj = color_elem.asMObject()
1177  num_data_fn = om.MFnNumericData(color_obj)
1178  new_color = num_data_fn.getData()
1179  # The stock shaders expect color as float4.
1180  new_colors.append( (new_color[0], new_color[1], new_color[2], 1.0) )
1181  self.shape_colors = new_colors
1182 
1183  # Check for changes to color values / num colors
1184  if len(old_colors) != len(new_colors):
1185  self.update_mask |= self.kUpdateShaders
1186  else:
1187  for (o,n) in zip(old_colors, new_colors):
1188  if not isclose_tuple(o, n):
1189  self.update_mask |= self.kUpdateShaders
1190  break
1191 
1192  # Check for changes to lead/active colors
1193  old_lead_color = self.lead_color
1194  old_active_color = self.active_color
1195  view = omui.M3dView()
1196  lead_color_index = cmds.displayColor('lead', q=True, active=True) - 1
1197  lead_color = view.colorAtIndex( lead_color_index, omui.M3dView.kActiveColors )
1198  self.lead_color = (lead_color.r, lead_color.g, lead_color.b, lead_color.a)
1199 
1200  active_color_index = cmds.displayColor('active', q=True, active=True) - 1
1201  active_color = view.colorAtIndex( active_color_index, omui.M3dView.kActiveColors )
1202  self.active_color = (active_color.r, active_color.g, active_color.b, active_color.a)
1203 
1204  old_sel_colors = (old_lead_color, old_active_color)
1205  new_sel_colors = (self.lead_color, self.active_color)
1206  for (o, n) in zip(old_sel_colors, new_sel_colors):
1207  if not isclose_tuple(o, n):
1208  self.update_mask |= self.kUpdateShaders
1209  break
1210 
1211  # Check drawOnTop changes
1212  draw_on_top_plug = om.MPlug(self.node.thisMObject(), helperShapeNode.aDrawOnTop)
1213  old_draw_on_top = self.draw_on_top
1214  self.draw_on_top = draw_on_top_plug.asBool()
1215  if self.draw_on_top != old_draw_on_top:
1216  self.update_mask |= self.kUpdateShaders
1217 
1218  def requires_update_geometry(self):
1219  '''
1220  Checks if any attributes have changed that would affect
1221  the geometry of any render items. If so, sets the kUpdateGeometry
1222  and kUpdateMatrix update bits.
1223  '''
1224  shape_plug = om.MPlug(self.node.thisMObject(), helperShapeNode.aShape)
1225  old_shape_index = self.shape_index
1226  self.shape_index = shape_plug.asInt()
1227  if self.shape_index >= len(helperShapeNode.shapes):
1228  self.shape_index = old_shape_index
1229  if not self.shape_index == old_shape_index:
1230  old_shape = helperShapeNode.shapes[old_shape_index]
1231  for item in old_shape.items:
1232  item.clear_buffers()
1233 
1234  # If geometry is requires update, then render items
1235  # will be rebuilt and matrices will also require
1236  # updates.
1237  self.update_mask |= self.kUpdateGeometry
1238  self.update_mask |= self.kUpdateMatrix
1239 
1240  def requires_update_matrix(self):
1241  '''
1242  Checks if any attributes have changed that would affect
1243  the matrices of any render items. If so, sets the
1244  kUpdateMatrix update bit.
1245  '''
1246  old_instances = self.instances
1247 
1248  # Update cached instance data.
1249  dag_fn = om.MFnDagNode(self.node.thisMObject())
1250  paths = dag_fn.getAllPaths()
1251  self.instances = {}
1252  for path in paths:
1253  if not path.isVisible():
1254  continue
1255  matrix = path.inclusiveMatrix()
1256  display_status = omr.MGeometryUtilities.displayStatus(path)
1257  self.instances[path.instanceNumber()] = (matrix, display_status)
1258 
1259  # Check for changes in total number of instances.
1260  old_instance_nums = set(old_instances.keys())
1261  new_instance_nums = set(self.instances.keys())
1262  instance_num_diff = old_instance_nums ^ new_instance_nums
1263  if len(instance_num_diff):
1264  # If the number of instances changes, rebuild all items.
1265  self.update_mask |= self.kUpdateGeometry
1266  self.update_mask |= self.kUpdateMatrix
1267  return
1268 
1269  # Check for changes in matrix or display_status.
1270  for i in self.instances.keys():
1271  (old_matrix, old_display) = old_instances[i]
1272  (new_matrix, new_display) = self.instances[i]
1273 
1274  if old_display != new_display:
1275  self.update_mask |= self.kUpdateMatrix
1276  return
1277  if not old_matrix.isEquivalent(new_matrix):
1278  self.update_mask |= self.kUpdateMatrix
1279  return
1280 
1281  def update(self, container, frame_context):
1282  '''
1283  Updates only the portions of the override that require
1284  updating.
1285  '''
1286  if self.update_mask & self.kUpdateShaders:
1287  self.update_shaders()
1288  if self.update_mask & self.kUpdateGeometry:
1289  self.update_geometry(container, frame_context)
1290  if self.update_mask & self.kUpdateMatrix:
1291  self.update_matrix(container, frame_context)
1292  if self.update_mask & self.kUpdateShaders and not self.update_mask & self.kUpdateGeometry:
1293  self.update_shaders_assign(container)
1294  self.update_mask = 0
1295 
1296  def update_shaders(self):
1297  '''
1298  Initialize shaders if uninitialized and set their
1299  parameters (ex. color)
1300  '''
1301  shader_mgr = omr.MRenderer.getShaderManager()
1302  if not shader_mgr:
1303  return
1304 
1305  if self.draw_on_top:
1306  pre_cb = helperShapeSubSceneOverride.pre_shader_cb
1307  post_cb = helperShapeSubSceneOverride.post_shader_cb
1308  else:
1309  pre_cb = None
1310  post_cb = None
1311 
1312  num_colors = len(self.shape_colors)
1313  self.wire_shaders = []
1314  self.shaded_shaders = []
1315  for i in xrange(num_colors):
1316  shader = self.shader_cache.get_wire_shader(self.shape_colors[i], self.draw_on_top, pre_cb, post_cb)
1317  self.wire_shaders.append(shader)
1318  for i in xrange(num_colors):
1319  shader = self.shader_cache.get_shaded_shader(self.shape_colors[i], self.draw_on_top, pre_cb, post_cb)
1320  self.shaded_shaders.append(shader)
1321  self.lead_select_shader = self.shader_cache.get_wire_shader(self.lead_color, self.draw_on_top, pre_cb, post_cb)
1322  self.active_select_shader = self.shader_cache.get_wire_shader(self.active_color, self.draw_on_top, pre_cb, post_cb)
1323 
1324  def update_shaders_assign(self, container):
1325  '''
1326  Update shader assignments on geometry items.
1327 
1328  This is required for the case when shaders require updating,
1329  but geometry does not. In these cases we must update the
1330  assigned shaders on the existing render items.
1331 
1332  A separate routine is required for assignment so that both
1333  shaders and render item lists are properly initialized prior
1334  to assignment.
1335  '''
1336  for (i, item_name) in self.wire_items:
1337  item = container.find(item_name)
1338  shader = self.wire_shaders[i]
1339  if not shader is item.getShader():
1340  item.setShader(shader)
1341  for (i, item_name) in self.shaded_items:
1342  item = container.find(item_name)
1343  shader = self.shaded_shaders[i]
1344  if not shader is item.getShader():
1345  item.setShader(shader)
1346  for (_, item_name) in self.lead_items:
1347  item = container.find(item_name)
1348  if not shader is item.getShader():
1349  item.setShader(self.lead_select_shader)
1350  for (_, item_name) in self.active_items:
1351  item = container.find(item_name)
1352  if not shader is item.getShader():
1353  item.setShader(self.active_select_shader)
1354 
1355  def update_geometry(self, container, frame_context):
1356  '''
1357  Update the subscene container with the render items to draw the shape.
1358  This method regenerates the render items from a clean slate.
1359  '''
1360  container.clear()
1361 
1362  # For later searches into the container, record the names
1363  # of each render item by its selection status.
1364  self.lead_items = []
1365  self.active_items = []
1366  self.wire_items = []
1367  self.shaded_items = []
1368 
1369  active_shape = helperShapeNode.shapes[self.shape_index]
1370  for (i,item) in enumerate(active_shape.items):
1371  bounds = item.bounding_box()
1372 
1373  # Shared render items (wire, lead, active)
1374  vertex_buffer = item.vertex_buffer()
1375  wire_index_buffer = item.wire_index_buffer()
1376 
1377  wire_name = '_'.join((active_shape.name, str(i), self.sWireName))
1378  wire_item = omr.MRenderItem.create( wire_name, omr.MRenderItem.NonMaterialSceneItem, omr.MGeometry.kLines )
1379  wire_item.setDrawMode(omr.MGeometry.kWireframe | omr.MGeometry.kShaded | omr.MGeometry.kTextured)
1380  if item.type == helperShapeRepItemType.kMesh:
1381  # Mesh wireframe only visible in wireframe mode.
1382  wire_item.setDrawMode(omr.MGeometry.kWireframe)
1383  wire_item.setDepthPriority(omr.MRenderItem.sDormantWireDepthPriority)
1384  wire_item.setShader(self.wire_shaders[i])
1385  container.add(wire_item)
1386  self.setGeometryForRenderItem(wire_item, vertex_buffer, wire_index_buffer, bounds)
1387  self.wire_items.append((i, wire_name))
1388 
1389  lead_name = '_'.join((active_shape.name, str(i), self.sLeadName))
1390  lead_item = omr.MRenderItem.create( lead_name, omr.MRenderItem.DecorationItem, omr.MGeometry.kLines )
1391  lead_item.setDrawMode(omr.MGeometry.kWireframe | omr.MGeometry.kShaded | omr.MGeometry.kTextured)
1392  lead_item.setDepthPriority(omr.MRenderItem.sActiveWireDepthPriority)
1393  lead_item.setShader(self.lead_select_shader)
1394  container.add(lead_item)
1395  self.setGeometryForRenderItem(lead_item, vertex_buffer, wire_index_buffer, bounds)
1396  self.lead_items.append((i, lead_name))
1397 
1398  active_name = '_'.join((active_shape.name, str(i), self.sActiveName))
1399  active_item = omr.MRenderItem.create( active_name, omr.MRenderItem.DecorationItem, omr.MGeometry.kLines )
1400  active_item.setDrawMode(omr.MGeometry.kWireframe | omr.MGeometry.kShaded | omr.MGeometry.kTextured)
1401  active_item.setDepthPriority(omr.MRenderItem.sActiveWireDepthPriority)
1402  active_item.setShader(self.active_select_shader)
1403  container.add(active_item)
1404  self.setGeometryForRenderItem(active_item, vertex_buffer, wire_index_buffer, bounds)
1405  self.active_items.append((i, active_name))
1406 
1407  # Type specific render items (shaded)
1408  if item.type == helperShapeRepItemType.kMesh:
1409  shaded_index_buffer = item.shaded_index_buffer()
1410  shaded_name = '_'.join((active_shape.name, str(i), self.sShadedName))
1411  shaded_item = omr.MRenderItem.create( shaded_name, omr.MRenderItem.NonMaterialSceneItem, omr.MGeometry.kTriangles )
1412  shaded_item.setDrawMode( omr.MGeometry.kShaded | omr.MGeometry.kTextured )
1413  shaded_item.setExcludedFromPostEffects(True)
1414  shaded_item.setCastsShadows(False)
1415  shaded_item.setReceivesShadows(False)
1416  shaded_item.setShader(self.shaded_shaders[i])
1417  container.add(shaded_item)
1418  self.setGeometryForRenderItem(shaded_item, vertex_buffer, shaded_index_buffer, bounds)
1419  self.shaded_items.append((i, shaded_name))
1420 
1421  def update_matrix(self, container, frame_context):
1422  '''
1423  Updates the matrices of the render items in the container.
1424  '''
1425  # Sort items by display status.
1426  num_instances = len(self.instances)
1427  lead_instances = om.MMatrixArray()
1428  active_instances = om.MMatrixArray()
1429  dormant_instances = om.MMatrixArray()
1430 
1431  for (_, instance) in self.instances.items():
1432  (matrix, display_status) = instance
1433  if display_status == omr.MGeometryUtilities.kLead:
1434  lead_instances.append(matrix)
1435  elif display_status == omr.MGeometryUtilities.kActive:
1436  active_instances.append(matrix)
1437  # Always display dormant items
1438  dormant_instances.append(matrix)
1439 
1440  dormant_items = []
1441  dormant_items.extend(self.wire_items)
1442  dormant_items.extend(self.shaded_items)
1443 
1444  if num_instances == 0:
1445  # Disable all render items
1446  #
1447  # BUG: The API docs and function signature suggest this
1448  # expects only the self argument, but it errors expecting
1449  # 1 argument. Pass in a dummy value of 0 as a workaround.
1450  container_it = container.getIterator(0)
1451  item = container_it.next()
1452  while item is not None:
1453  item.enable(False)
1454  item = container_it.next()
1455  container_it.destroy()
1456  container_it = None
1457  elif num_instances == 1:
1458  # Set our render items directly.
1459  matrix = dormant_instances[0]
1460  for (_, item_name) in dormant_items:
1461  item = container.find(item_name)
1462  item.enable(True)
1463  item.setMatrix(matrix)
1464  if WANT_CONSOLIDATION:
1465  item.setWantSubSceneConsolidation(True)
1466  for (_, item_name) in self.lead_items:
1467  item = container.find(item_name)
1468  item.enable(len(lead_instances) > 0)
1469  item.setMatrix(matrix)
1470  if WANT_CONSOLIDATION:
1471  item.setWantSubSceneConsolidation(True)
1472  for (_, item_name) in self.active_items:
1473  item = container.find(item_name)
1474  item.enable(len(active_instances) > 0)
1475  item.setMatrix(matrix)
1476  if WANT_CONSOLIDATION:
1477  item.setWantSubSceneConsolidation(True)
1478  else:
1479  # Set instance transforms
1480  for (_, item_name) in dormant_items:
1481  item = container.find(item_name)
1482  item.enable(True)
1483  self.setInstanceTransformArray(item, dormant_instances)
1484  has_lead_instances = len(lead_instances) > 0
1485  for (_, item_name) in self.lead_items:
1486  item = container.find(item_name)
1487  item.enable(has_lead_instances)
1488  if has_lead_instances:
1489  self.setInstanceTransformArray(item, lead_instances)
1490  has_active_instances = len(active_instances) > 0
1491  for (_, item_name) in self.active_items:
1492  item = container.find(item_name)
1493  item.enable(has_active_instances)
1494  if has_active_instances:
1495  self.setInstanceTransformArray(item, active_instances)
1496 
1497 
1498 ###############################################################################
1499 #
1500 # helperShape plug-in initialize/uninitialize
1501 #
1502 def initializePlugin(obj):
1503  plugin = om.MFnPlugin( obj )
1504 
1505  plugin_path = os.path.normpath(plugin.loadPath())
1506  lib_path = os.path.join(plugin_path, 'library')
1507  lib_env_var = 'MAYA_LOCATOR_HELPER_SHAPE_LIB'
1508  if lib_env_var in os.environ:
1509  env_path = os.environ[lib_env_var]
1510  lib_path = os.path.normpath(env_path)
1511  helperShapeNode.init_shapes(lib_path)
1512 
1513  try:
1514  plugin.registerNode(helperShapeNode.name,
1515  helperShapeNode.id,
1516  helperShapeNode.creator,
1517  helperShapeNode.initialize,
1518  om.MPxNode.kLocatorNode,
1519  helperShapeNode.drawDbClassification)
1520  except:
1521  sys.stderr.write('Failed to register locatorHelperShape node.\n')
1522  raise
1523 
1524  try:
1525  mel.eval(helperShapeNode.attrEditorTemplate)
1526  except:
1527  sys.stderr.write('Failed to load locatorHelperShape AETemplate script.\n')
1528  raise
1529 
1530  try:
1531  omr.MDrawRegistry.registerSubSceneOverrideCreator(
1532  helperShapeNode.drawDbClassification,
1533  helperShapeNode.drawRegistrantId,
1534  helperShapeSubSceneOverride.creator)
1535  except:
1536  sys.stderr.write('Failed to register locatorHelperShape SubSceneOverride.\n')
1537  raise
1538 
1539  try:
1540  plugin.registerCommand(helperShapeExportCmd.name,
1541  helperShapeExportCmd.creator,
1542  helperShapeExportCmd.new_syntax)
1543  except:
1544  sys.stderr.write('Failed to register locatorHelperShapeExportCmd.\n')
1545  raise
1546 
1547  # Register helperShape selection mask
1548  om.MSelectionMask.registerSelectionType(HELPER_SHAPE_SELECTION_MASK, HELPER_SHAPE_SELECTION_PRIORITY)
1549  cmds.selectType(byName=(HELPER_SHAPE_SELECTION_MASK, True))
1550 
1551 def uninitializePlugin(obj):
1552  plugin = om.MFnPlugin( obj )
1553 
1554  try:
1555  plugin.deregisterNode(helperShapeNode.id)
1556  except:
1557  sys.stderr.write('Failed to deregister locatorHelperShape node.\n')
1558  pass
1559 
1560  try:
1561  omr.MDrawRegistry.deregisterSubSceneOverrideCreator(
1562  helperShapeNode.drawDbClassification,
1563  helperShapeNode.drawRegistrantId)
1564  except:
1565  sys.stderr.write('Failed to deregister locatorHelperShape SubSceneOverride.\n')
1566  pass
1567 
1568  try:
1569  plugin.deregisterCommand(helperShapeExportCmd.name)
1570  except:
1571  sys.stderr.write('Failed to deregister locatorHelperShapeExportCmd.\n')
1572  pass
1573 
1574  # Deregister helperShape selection mask
1575  om.MSelectionMask.deregisterSelectionType(HELPER_SHAPE_SELECTION_MASK)
1576