VisualiserTool : Add primitive variable menu

Changes.md

@@ -1,6 +1,11 @@
 1.5.x.x (relative to 1.5.2.0)
 =======
 
+Improvements
+------------
+
+- VisualiserTool : Changed `dataName` input widget for choosing the primitive variable to visualise to a list of available variable names for the current selection.
+
 Fixes
 -----
 

python/GafferSceneUI/VisualiserToolUI.py

@@ -34,10 +34,17 @@
 #
 ##########################################################################
 
+import functools
+
+import IECore
+import IECoreScene
+
 import Gaffer
 import GafferUI
 import GafferSceneUI
 
+from GafferUI.PlugValueWidget import sole
+
 Gaffer.Metadata.registerNode(
 
 	GafferSceneUI.VisualiserTool,
@@ -58,13 +65,15 @@
 
 			"description",
 			"""
-			Specifies the name of the primitive variable to visualise. The data should
-			be of type float, V2f or V3f.
+			Specifies the name of the primitive variable to visualise. Variables of
+			type int, float, V2f, Color3f or V3f can be visualised.
 			""",
 
 			"toolbarLayout:section", "Bottom",
 			"toolbarLayout:width", 150,
 
+			"plugValueWidget:type", "GafferSceneUI.VisualiserToolUI._DataNameChooser",
+
 		],
 		"opacity" : [
 
@@ -118,3 +127,86 @@
 
 	},
 )
+
+class _DataNameChooser( GafferUI.PlugValueWidget ) :
+
+	def __init__( self, plug, **kw ) :
+
+		self.__menuButton = GafferUI.MenuButton(
+			text = plug.getValue(),
+			menu = GafferUI.Menu( Gaffer.WeakMethod( self.__menuDefinition ) )
+		)
+
+		GafferUI.PlugValueWidget.__init__( self, self.__menuButton, plug, **kw )
+
+	def _updateFromValues( self, values, exception ) :
+
+		self.__menuButton.setText( sole( values ) or "None" )
+
+	def __menuDefinition( self ) :
+
+		menuDefinition = IECore.MenuDefinition()
+
+		node = self.getPlug().node()
+		if not isinstance( node, GafferSceneUI.VisualiserTool ) :
+			return
+		if self.getPlug() != node["dataName"] :
+			return
+
+		scenePlug = node.view()["in"].getInput()
+		scriptNode = node.view().scriptNode()
+		with node.view().context() :
+			selection = GafferSceneUI.ScriptNodeAlgo.getSelectedPaths( scriptNode )
+
+			primVars = set()
+
+			for path in selection.paths() :
+				if not scenePlug.exists( path ) :
+					continue
+
+				primitive = scenePlug.object( path )
+				if not isinstance( primitive, IECoreScene.MeshPrimitive ) :
+					continue
+
+				for v in primitive.keys() :
+					if primitive[v].interpolation not in [
+						IECoreScene.PrimitiveVariable.Interpolation.FaceVarying,
+						IECoreScene.PrimitiveVariable.Interpolation.Uniform,
+						IECoreScene.PrimitiveVariable.Interpolation.Vertex,
+					] :
+						continue
+
+					if not isinstance(
+						primitive[v].data,
+						(
+							IECore.IntVectorData,
+							IECore.FloatVectorData,
+							IECore.V2fVectorData,
+							IECore.Color3fVectorData,
+							IECore.V3fVectorData,
+						)
+					) :
+						continue
+
+					primVars.add( v )
+
+		if len( primVars ) == 0 :
+			menuDefinition.append( "/None Available", { "active" : False } )
+
+		else :
+			for v in reversed( sorted( primVars ) ) :
+				menuDefinition.prepend(
+					"/" + v,
+					{
+						"command" : functools.partial( Gaffer.WeakMethod( self.__setDataName ), v ),
+						"checkBox" : self.getPlug().getValue() == v,
+					}
+				)
+
+		menuDefinition.prepend( "/PrimVarDivider", { "divider" : True, "label" : "Primitive Variables" } )
+
+		return menuDefinition
+
+	def __setDataName( self, value, *unused ) :
+
+		self.getPlug().setValue( value )