SelectionToolUI : Custom dropdown menu

python/GafferSceneUI/SelectionToolUI.py

@@ -34,6 +34,7 @@
 #
 ##########################################################################
 
+import functools
 import imath
 
 import IECore
@@ -82,14 +83,7 @@
 			which may differ from what is originally selected.
 			""",
 
-			"plugValueWidget:type", "GafferUI.PresetsPlugValueWidget",
-
-			"presetNames", lambda plug : IECore.StringVectorData(
-				GafferSceneUI.SelectionTool.registeredSelectModeLabels()
-			),
-			"presetValues", lambda plug : IECore.StringVectorData(
-				GafferSceneUI.SelectionTool.registeredSelectModes()
-			),
+			"plugValueWidget:type", "GafferSceneUI.SelectionToolUI.SelectModePlugValueWidget",
 
 			"toolbarLayout:section", "Bottom",
 			"toolbarLayout:width", 150,
@@ -110,3 +104,77 @@ class _RightSpacer( GafferUI.Spacer ) :
 	def __init__( self, imageView, **kw ) :
 
 		GafferUI.Spacer.__init__( self, size = imath.V2i( 0, 0 ) )
+
+class SelectModePlugValueWidget( GafferUI.PlugValueWidget ) :
+
+	def __init__( self, plugs, **kw ) :
+
+		self.__menuButton = GafferUI.MenuButton( "", menu = GafferUI.Menu( Gaffer.WeakMethod( self.__menuDefinition ) ) )
+		GafferUI.PlugValueWidget.__init__( self, self.__menuButton, plugs, **kw )
+
+	def _updateFromValues( self, values, exception ) :
+
+		if exception is not None :
+			self.__menuButton.setText( "" )
+		else :
+			modifiers = dict(
+				zip(
+					GafferSceneUI.SelectionTool.registeredSelectModes(),
+					GafferSceneUI.SelectionTool.registeredSelectModeLabels()
+				)
+			)
+
+			assert( len( values ) == 1 )
+			label = modifiers.get( values[0], None )
+
+			if label is not None :
+				self.__menuButton.setText( "Any" if label == "Any" else label.partition( "/" )[-1] )
+			else :
+				self.__menuButton.setText( "Invalid" )
+
+		self.__menuButton.setErrored( exception is not None )
+
+	def _updateFromEditable( self ) :
+
+		self.__menuButton.setEnabled( self._editable() )
+
+	def __menuDefinition( self ) :
+
+		result = IECore.MenuDefinition()
+
+		modifiers = dict(
+			zip(
+				GafferSceneUI.SelectionTool.registeredSelectModes(),
+				GafferSceneUI.SelectionTool.registeredSelectModeLabels()
+			)
+		)
+
+		modifierCategories = set()
+
+		with self.getContext() :
+			currentValue = self.getPlug().getValue()
+
+		for modifierName, modifierLabel in modifiers.items() :
+			category, sep, label = modifierLabel.partition( "/" )
+
+			if category == "Any" :
+				label = "Any"
+
+			if category != "Any" and category not in modifierCategories :
+				result.append( f"/__{category}Dividier", { "divider" : True, "label" : category } )
+				modifierCategories.add( category )
+
+			result.append(
+				f"/{label}",
+				{
+					"command" : functools.partial( Gaffer.WeakMethod( self.__setValue ), modifierName ),
+					"checkBox" : modifierName == currentValue
+				}
+			)
+
+		return result
+
+	def __setValue( self, modifier, *unused ) :
+
+		with Gaffer.UndoScope( self.getPlug().ancestor( Gaffer.ScriptNode ) ) :
+			self.getPlug().setValue( modifier )