toggle_weight_paint.py

@@ -1,19 +1,25 @@
 import bpy
 
-# This operator is to make entering weight paint mode less of a pain.
-# You just need to select a mesh, run the operator, and you should be ready to weight paint.
+# This operator is added to the Object menu.
 
-# It registers an operator called "Toggle Weight Paint Mode" that does the following:
+# It does the following:
 #	Set active object to weight paint mode
+#	Make sure active object is not in wireframe or bounding box display type
 #	Set shading mode to a Flat, Single Color, White shading
 #	Find first armature via the object's modifiers.
-#	Ensure it is visible by moving it to a temporary collection and changing its visibility settings.
-#	Enable "In Front" option
-#	Set it to pose mode.
-# When running the operator again, it should restore all modes and shading settings, delete the temporary collection, and restore the armature's visibility settings.
-# You need to set up your own keybind for this operator.
+#	Ensure it is visible, select it and set it to pose mode.
 
-coll_name = "Weight Paint"
+# This allows you to start weight painting with a single button press from any state.
+
+# When running the operator again, it should restore everything to how it was before.
+
+coll_name = "Weight Paint Toggle Helper"
+
+def get_armature_of_meshob(obj: bpy.types.Object):
+	"""Find and return the armature that deforms this mesh object."""
+	for m in obj.modifiers:
+		if m.type=='ARMATURE':
+			return m.object
 
 class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 	"""Toggle weight paint mode properly with a single operator. """
@@ -21,51 +27,75 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 	bl_label = "Toggle Weight Paint Mode"
 	bl_options = {'REGISTER', 'UNDO'}
 
-	# local_view: bpy.props.BoolProperty(name="Local View", description="Enter Local view with the mesh and armature")
+	# local_view: bpy.props.BoolProperty(
+	# 	name		 = "Local View"
+	# 	,description = "Enter Local view with the mesh and armature"
+	# )
 	
 	@classmethod
 	def poll(cls, context):
-		return context.object and context.object.type=='MESH'
+		ob = context.object
+		return ob and ob.type=='MESH'
 
 	def draw(self, context):
 		self.layout.operator(EASYWEIGHT_OT_toggle_weight_paint.bl_idname)
 
-	def execute(self, context):
+	def enter_wp(self, context):
+		"""Enter weight paint mode, change the necessary settings, and save their
+		original states so they can be restored when leaving wp mode."""
+
 		obj = context.object
 
-		mode = obj.mode
-		enter_wp = not (mode == 'WEIGHT_PAINT')
-
-		# Finding armature.
-		armature = None
-		for m in obj.modifiers:
-			if m.type=='ARMATURE':
-				armature = m.object
-
-		if enter_wp:
-			### Entering weight paint mode. ###
-
-			# If the mesh object's display mode was anything other than Solid or Textured, store it, as we'll have to change it. and then change it back.
-			if obj.display_type not in ['SOLID', 'TEXTURED']:
+		# Store old display_type setting in a Custom Property on the Object.
 		obj['wpt_display_type'] = obj.display_type
+		# Ensure display_type is SOLID; weights don't display in WIRE or BOUNDS.
+		if obj.display_type not in ['SOLID', 'TEXTURED']:
 			obj.display_type = 'SOLID'
 
-			obj.data.use_mirror_topology = False
-
-			# Store old shading settings in a dict custom property
+		# Store old shading settings in a Custom Property dictionary in the Scene.
 		if 'wpt' not in context.screen:
 			context.screen['wpt'] = {}
+
 		wpt = context.screen['wpt'].to_dict()
-			# Set modes.
-			if armature:
+
+		# If we are entering WP mode for the first time or if the last time 
+		# the operator was exiting WP mode, then save current shading info.
+		if 'last_switch_in' not in wpt or wpt['last_switch_in']==False:
+			context.screen['wpt']['shading_type'] = context.space_data.shading.type
+			if context.space_data.shading.type=='SOLID':
+				# These properties only exist when the shading type is SOLID.
+				context.screen['wpt']['light'] = context.space_data.shading.light
+				context.screen['wpt']['color_type'] = context.space_data.shading.color_type
+				context.screen['wpt']['single_color'] = context.space_data.shading.single_color
+			context.screen['wpt']['active_object'] = obj
+
+		# This flag indicates that the last time this operator ran, we were 
+		# switching INTO wp mode.
+		context.screen['wpt']['last_switch_in'] = True
+		context.screen['wpt']['mode'] = obj.mode
+
+		# Set shading
+		if context.space_data.shading.type=='SOLID':
+			context.space_data.shading.light = 'FLAT'
+			context.space_data.shading.color_type = 'SINGLE'
+			context.space_data.shading.single_color = (1,1,1)
+
+		# Enter WP mode.
+		bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
+
+		# If there is an armature, make sure it is visible.
+		armature = get_armature_of_meshob(obj)
+		if not armature:
+			return {'FINISHED'}
 		context.screen['wpt']['armature_enabled'] = armature.hide_viewport
 		context.screen['wpt']['armature_hide'] = armature.hide_get()
 		context.screen['wpt']['armature_in_front'] = armature.show_in_front
 		armature.hide_viewport = False
 		armature.hide_set(False)
 		armature.show_in_front = True
-				context.view_layer.objects.active = armature
 
+		# If the armature is still not visible, use a collection.
+		if not armature.visible_get():
 			coll = bpy.data.collections.get(coll_name)
 			if not coll:
 				coll = bpy.data.collections.new(coll_name)
@@ -74,64 +104,67 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 			
 			if armature.name not in coll.objects:
 				coll.objects.link(armature)
+
 		if not armature.visible_get():
-					print("By some miracle, the armature is still hidden, cannot reveal it for weight paint mode.")
-					armature=None
+			self.report({'WARNING'}, "Failed to make the armature visible. Curse you, Blender!!!")
 		else:
+			context.view_layer.objects.active = armature
 			bpy.ops.object.mode_set(mode='POSE')
+		
 		context.view_layer.objects.active = obj
-			bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
-			if 'last_switch_in' not in wpt or wpt['last_switch_in']==False:	# Only save shading info if we exitted weight paint mode using this operator.
-				context.screen['wpt']['shading_type'] = context.space_data.shading.type
-				if context.space_data.shading.type=='SOLID':
-					context.screen['wpt']['light'] = context.space_data.shading.light
-					context.screen['wpt']['color_type'] = context.space_data.shading.color_type
-					context.screen['wpt']['single_color'] = context.space_data.shading.single_color
-				context.screen['wpt']['active_object'] = obj
+		return {'FINISHED'}
 
-			context.screen['wpt']['last_switch_in'] = True	# Store whether the last time the operator ran, were we switching into or out of weight paint mode.
-			context.screen['wpt']['mode'] = mode
+	def leave_wp(self, context):
+		"""Leave weight paint mode, then find, restore, and delete the data
+		that was stored about shading settings in enter_wp()."""
 
-			# Set shading
-			if context.space_data.shading.type=='SOLID':
-				context.space_data.shading.light = 'FLAT'
-				context.space_data.shading.color_type = 'SINGLE'
-				context.space_data.shading.single_color = (1,1,1)
-
-		else:
-			### Leaving weight paint mode. ###
+		obj = context.object
 
 		# Restore object display type
 		if 'wpt_display_type' in obj:
 			obj.display_type = obj['wpt_display_type']
 			del obj['wpt_display_type']
 
-			if 'wpt' in context.screen:
-				info = context.screen['wpt'].to_dict()
+		if 'wpt' not in context.screen:
+			# There is no saved data to restore from, nothing else to do.
+			bpy.ops.object.mode_set(mode='OBJECT')
+			return {'FINISHED'}
+
+		wpt = context.screen['wpt'].to_dict()
+
 		# Restore mode.
-				bpy.ops.object.mode_set(mode=info['mode'])
+		bpy.ops.object.mode_set(mode=wpt['mode'])
+
+		# Flag to save that the last time the operator ran we were EXITING wp mode.
 		context.screen['wpt']['last_switch_in'] = False
 
 		# Restore shading options.
-				context.space_data.shading.type = info['shading_type']
+		context.space_data.shading.type = wpt['shading_type']
 		if context.space_data.shading.type=='SOLID':
-					context.space_data.shading.light = info['light']
-					context.space_data.shading.color_type = info['color_type']
-					context.space_data.shading.single_color = info['single_color']
+			context.space_data.shading.light = wpt['light']
+			context.space_data.shading.color_type = wpt['color_type']
+			context.space_data.shading.single_color = wpt['single_color']
 
-				# If the armature was un-hidden, hide it again.
-				if armature:
-					armature.hide_viewport = info['armature_enabled']
-					armature.hide_set(info['armature_hide'])
-					armature.show_in_front = info['armature_in_front']
+		armature = get_armature_of_meshob(obj)
+		if not armature:
+			return {'FINISHED'}
+		# If an armature was un-hidden, hide it again.
+		armature.hide_viewport = wpt['armature_enabled']
+		armature.hide_set(wpt['armature_hide'])
+		armature.show_in_front = wpt['armature_in_front']
 		coll = bpy.data.collections.get(coll_name)
 		if coll:
 			bpy.data.collections.remove(coll)
-			else:
-				# If we didn't enter weight paint mode with this operator, just go into object mode when trying to leave WP mode with this operator.
-				bpy.ops.object.mode_set(mode='OBJECT')
+		return {'FINISHED'}
+
+	def execute(self, context):
+		obj = context.object
+
+		if obj.mode != 'WEIGHT_PAINT':
+			return self.enter_wp(context)
+		else:
+			return self.leave_wp(context)
 
-		return { 'FINISHED' }
 
 def register():
 	from bpy.utils import register_class